Graph traversal of torch.nn.Module

nscotto: <p>Hi, I already asked a question about it but after coming back to it I think the answer is missing some points.<br>
So I would like to traverse the graph of modules, starting from the input in a depth/breadth first search manner.<br>
My goal is to automatically search &amp; replace blocks (subgraphs) in a model (graph).</p>
<p>Here is my <a href="https://discuss.pytorch.org/t/graph-exploration-of-unknown-models/103404/4">previous answered question</a></p>
<p>So after looking at the doc it seems that <code>torch.nn.Module.named_modules()</code> iterates over all <code>torch.nn.Module</code> found in the constructor, I imagine respecting the order in which they are declared. The problem is that it does not necessarily follow a depth first/breadth first, even more if the network doesn’t follow a sequential structure.<br>
The computation graph is defined by the order of computation in the forward method, not the order of declaration in the constructor…<br>
Another (“more minor”) problem is that when using functional api for some computation (like using <code>torch.nn.functional.relu()</code> in the <code>forward</code> method), <code>named_modules()</code> can’t see it.</p>
<p>So I am still looking for a way to do a depth/breadth first graph traversal in order to operate search and replace of blocks.<br>
Is there a way I can achieve this?</p>

albanD: <p>Hi,</p>
<p>One thing you most likely want to checkout is <a href="https://pytorch.org/docs/master/fx.html?highlight=fx#module-torch.fx" class="inline-onebox">torch.fx — PyTorch master documentation</a></p>

nscotto: <p>It looks promising, and will probably help me achieving what I want to do right now.<br>
Any idea on when it could be released?</p>

albanD: <p>It will be a beta feature in the upcoming 1.8 release.</p>

nscotto: <p>Right, I imagine you don’t want to adventure into release date forecast <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=9" title=":sweat_smile:" class="emoji" alt=":sweat_smile:"><br>
Can’t wait for it! <img src="https://discuss.pytorch.org/images/emoji/apple/grinning.png?v=9" title=":grinning:" class="emoji" alt=":grinning:"></p>
<p>Cheers</p>

albanD: <aside class="quote no-group" data-username="nscotto" data-post="5" data-topic="110821">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/nscotto/48/29758_2.png" class="avatar"> nscotto:</div>
<blockquote>
<p>Right, I imagine you don’t want to adventure into release date forecast <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=12" title=":sweat_smile:" class="emoji" alt=":sweat_smile:" loading="lazy" width="20" height="20"></p>
</blockquote>
</aside>
<p>The release branch is planned to be cut later today.<br>
So between testing, release notes and all, you should have it at the end of the month <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

nscotto: <p>Wow, good work guys that’s awesome!</p>