FX: insert a graph instead of a node

torch.compile FX (Functional Transformations)
1

the fx api has methods for inserting nodes, it is flexible enough so that the node I insert can be a fully featured model with several layers, however I am facing a problem in which I need to insert each nodes from the submodel instead of the submodel as one node and I am not sure how to achieve that.

The context is that I am writing a backbone replacer, it automatically detects backbones and can replace them. It is fine to just remove all nodes from the old backbone and insert the new backbone as a single node in most cases, but in some cases there are skip connections between the backbone and the head and I would like to be able to handle that case. That means I need to insert the replacing backbone layer by layer so I can remap the previous skip connections to layers of the replacing backbone.

I hope that made sense!
Thank you in advance.

2

What I would somewhat like to do is:

model_nodes = list(model.graph.nodes)
block_nodes = list(model.graph.nodes)
block_nodes[0].replace_all_uses_with(model_nodes[insert_idx])
model_nodes[insert_idx].replace_all_uses_with(block_nodes[-2])

However I don’t think this works, the block nodes are referring to modules not present in the model. I can add the block as a module to the model, but I think replace_all_uses_with is only intended to work with nodes already added to the graph if that makes sense.

I would be really thankful to anyone willing to share some help.

3

So, I think there are a couple of different questions here:

  1. The first one is that you want to insert a module in op by op instead of the entire module at once. I think the easiest way of doing this is through creating proxies and implicitly creating your graph. You would use the inserting_before, and then initialize the proxies and call your module: examples/proxy_based_graph_creation.py at master · pytorch/examples · GitHub

  2. The other issue is that the new graph is referring to modules that don’t currently exist in your graph. You currently need to resolve this yourself, by inserting your modules into your GraphModule.

4

Thanks for the answer.
For the second point, how can I get a fresh instance of the module I want to insert if it is already in a model?
Would inserting a copy of it just work?