Extract graph for editing

torch.compile FX (Functional Transformations)
1

What I need is:

  1. to edit graph for original nn.Module model;
  2. train it
  3. to edit graph again

The only way I’ve found to achieve graph editing using pytorch2.0 tools is using custom_backend but the point that it triggers only one time so that my third case in the list above is not achievable.
What are the existing options?

2

This tutorial might help Google Colab, you can move from an aten graph back to an nn module easily with some of this infra

3
  1. is this applicable for aten graph only (not for torch graph)?
  2. if yes for first question than considering some example from the notebook
import torch._dynamo
from torch._functorch.aot_autograd import aot_module_simplified

def toy_backend(gm, sample_inputs): 
    def my_compiler(gm, sample_inputs):
        # <implement your compiler here>
        #place1
        gm.print_readable()
        return gm.forward

    # Invoke AOTAutograd
    return aot_module_simplified(
        gm,
        sample_inputs,
        fw_compiler=my_compiler
    )

torch._dynamo.reset()
fn = torch.compile(backend=toy_backend, dynamic=True)(model)

# triggers compilation of forward graph on the first run
out = fn(input)

# triggers compilation of backward graph on the first run
out.sum().backward()

we can change graph in the first time in place1 but what about second time?