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

username: crissallan
cooked: <p>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</p>
<p>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.</p>
<p>for example for a model defined like below:</p>
<pre data-code-wrap="python"><code class="lang-python">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)
</code></pre>
<p>I want a graph that shows the how the input is processed between “MyLayer” instance like below:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/a/faf4fff879f8cb10a64583af35e4187baafebba7.png" data-download-href="https://discuss.pytorch.org/uploads/default/faf4fff879f8cb10a64583af35e4187baafebba7" title="graph_example"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/a/faf4fff879f8cb10a64583af35e4187baafebba7_2_288x500.png" alt="graph_example" data-base62-sha1="zO4hcv17LNBbrPWpVug4pbC63LF" width="288" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/a/faf4fff879f8cb10a64583af35e4187baafebba7_2_288x500.png, https://discuss.pytorch.org/uploads/default/original/3X/f/a/faf4fff879f8cb10a64583af35e4187baafebba7.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/f/a/faf4fff879f8cb10a64583af35e4187baafebba7.png 2x" data-dominant-color="E8E9E9"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">graph_example</span><span class="informations">377×653 16.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Is there any library or method I can achieve this?</p>

username: srishti-git1110
cooked: <p>Hey, you might be looking for <a href="https://pypi.org/project/torchviz/" rel="noopener nofollow ugc">torchviz</a>, and make_dot therein.</p>

username: crissallan
cooked: <p>Thanks for replying,.</p>
<p>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</p>