graphutils supports torch17

[chicm-ms]

something changed in pytorch v1.6 which breaks graphutils:

1. torch.onnx.set_training is removed
2. in traced graph, some prim::Constant nodes can be shared as inputs of multiple nodes in torch1.6:
   for example:

import torch
import torch.nn as nn
import torch.nn.functional as F


class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.bn1 = nn.BatchNorm2d(self.conv1.out_channels)
        self.bn2 = nn.BatchNorm2d(self.conv2.out_channels)
        self.fc1 = nn.Linear(4 * 4 * 50, 10)

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.max_pool2d(x, 2, 2)
        x = x.view(x.size(0), -1)
        
        x = self.fc1(x)
        return x

class MyModel2(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.submodel = MyModel()

    def forward(self, x):
        return self.submodel(x)

In traced graph of MyModel2, we found that, conv1 and conv2 share a few prim::Constant input nodes, bn1 and bn2 also shares a few prim::Constant nodes as input, this breaks graphutils.

[QuanluZhang]

the element in output_to_node is a list? This means one output is generated by more than one node?

[chicm-ms]

assertion added to ensure the list has 1 element at most.

[liuzhe-lz]

Could you check if the test cases work with PyTorch 1.6+ in following scripts?

• test_graph_utils.py
• test_compression_utils.py
• test_pruners.py
• test_dependecy_aware.py

[chicm-ms]

pytorch 1.6+ turned on for

test_compression_utils.py
test_pruners.py
test_dependecy_aware.py

for test_graph_utils.py, there are some expected file are generated by Yuge, they may need to be upgraded for pytorch 1.6+

@ultmaster , would you please help upgrade the protobuf test cases in test_graph_utils.py for pytroch1.6+ ?

[zheng-ningxin]

If predecessor_node is already merged into this node group, then I guess it's output should not be taken as the input of this node group?

[chicm-ms]

checking added, please review again.

[zheng-ningxin]

I guess I didn't explain myself clearly, sorry about that.
What I want to say is that maybe we can remove the else part of the following code.

 if not self._is_key_func(predecessor_node):
    node_group.append(predecessor_node)
    node_queue.put(predecessor_node)
 else:
    inputs.add(input_name)

Because whether the input_name should be added to the inputs(inputs) of this nodes group is only decided by if predecessor_node is in nodes? It has nothing to do with whether predecessor_node is a key node/func? I guess? Please correct me if I'm wrong. Thanks~

[chicm-ms]

It does not work removing the else as I tested. Because if the predecessor node is key func, the predecessor node will be merged into another node. There is only one key func node within one merged node.

[zheng-ningxin]

Got it! Thanks~