How to extract a model graph's based on using computation graph?

autograd
1

Hey, I’d like to know how can I draw/extract a model graph. I know that computation graph can give me some graph-like results, but each nodes can only be some tensor output by a specific differentible operation

In my case, I want the node of the graph to be the layer of the model instead of a tensor. Which means I just want to know the flow of the tensor from some layer to some other layer.

for example for a model defined like below:

class MyLayer(torch.nn.Module):

    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.linear = torch.nn.Linear(in_dim, out_dim)
        self.relu = torch.nn.ReLU()
   def forward(x):
        return self.relu(self.linear(x))
      

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.my_layer1 = MyLayer(10, 5)
        self.my_layer2 = MyLayer(10, 5)
        self.my_layer3 = MyLayer(5, 3)


    def forward(self, x):
        x1 = self.my_layer1(x)
        x2 = self.my_layer2(x)
        x3 = x1 + x2
        return self.my_layer3(x3)

I want a graph that shows the how the input is processed between “MyLayer” instance like below:

graph_example
graph_example377×653 16.4 KB

Is there any library or method I can achieve this?

2

Hey, you might be looking for torchviz, and make_dot therein.

3

Thanks for replying,.

I had a try on this lib before. I would say it only solved part of my problem. e.g. If I defined some nested operations within some customized nn.Module class, torchviz can only give me the fine-grained nodes of the comutation graph. But I just need the whole custimized nn.Module instance as a single node in the graph. Maybe, I need to do some modification based on the torchviz to achieve my goal