Setting parent and children of nodes with custom params

username: marksaroufim
cooked: <p>Hi folks,</p>
<p>I’l looking to take a graph that looks like the below (left side) and turn it into the right side graph to make it fit on a device with small on device memory</p>
<p>I was wondering if there was a recommended way to do such transformations? My output graph is not a linkedlist of nodes as in a -&gt;b → c so I don’t think the append and prepend operations will help my use case. Basically I want to be able to manually set the parent and children of a node, set some parameters for that node/layer and then add a new cat operation and generate a new graph.</p>
<p>Adding a new cat operations seems to be simple: <a href="https://github.com/pytorch/pytorch/blob/master/torch/fx/experimental/merge_matmul.py#L189" class="inline-onebox">pytorch/torch/fx/experimental/merge_matmul.py at main · pytorch/pytorch · GitHub</a><br>
Returning a new graph also simple: <a href="https://github.com/pytorch/pytorch/blob/master/torch/fx/experimental/merge_matmul.py#L220" class="inline-onebox">pytorch/torch/fx/experimental/merge_matmul.py at main · pytorch/pytorch · GitHub</a><br>
Set parameters: Should I manually just write to the modules dictionary?<br>
Unsure how to set parent and children for nodes, have a node potentially have multiple children</p>
<p>Are any of the public examples suitable as a starting point? <a class="mention" href="/u/chillee">@Chillee</a> <a class="mention" href="/u/james_reed">@James_Reed</a></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/7/97b1097ed222545e39208d65cda81c0c3910ad04.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/97b1097ed222545e39208d65cda81c0c3910ad04" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/7/97b1097ed222545e39208d65cda81c0c3910ad04_2_690x425.jpeg" alt="image" data-base62-sha1="lDVoGTJZEaH0XaYYDxriMLAhV9q" width="690" height="425" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/7/97b1097ed222545e39208d65cda81c0c3910ad04_2_690x425.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/9/7/97b1097ed222545e39208d65cda81c0c3910ad04_2_1035x637.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/9/7/97b1097ed222545e39208d65cda81c0c3910ad04_2_1380x850.jpeg 2x" data-dominant-color="F9F9F9"><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">image</span><span class="informations">2048×1262 123 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>

username: Chillee
cooked: <p>There is no direct concept of a “parent” or “child” in FX - the closest analogues are “args” (i.e. what nodes your current node is using) and “users” (i.e. what nodes are using your current node). Furthermore, “users” is a property derived from “args”, so the only thing you should be modifying is what the args to your function are.</p>
<p>That being said, it’s not totally obvious to me what the graph rewrite you want to do is (where does conv3 go?), but I’d imagine it looks a bit like this (assuming you want to replace conv3 with concat?):</p>
<ol>
<li>create conv2_1 and conv2_2 from conv2 (probably using nn.Module.add_module)</li>
<li>create 2 call_module nodes that take in conv1 as an arg using <code>fx.Graph.call_module</code> (inserting in front of conv3)</li>
<li>create a concat node with conv2_1 and conv2_2 as input</li>
<li>Replace all uses of conv3 with concat using <code>fx.Node.replace_all_uses_with</code>
</li>
</ol>
<p>It might also be easier for you to use the Subgraph Rewriter API: <a href="https://pytorch.org/docs/stable/fx.html#subgraph-rewriting-with-replace-pattern" class="inline-onebox" rel="noopener nofollow ugc">torch.fx — PyTorch 1.8.1 documentation</a></p>