How can compiled model generates a torch.fx.GraphModule

torch.compile
1

Compile a model using optimized = optimize("inductor")(model), the optimized is of type <class 'torch._dynamo.eval_frame.OptimizedModule'>, how to convert it to torch.fx.GraphModule?

2

You shouldn’t need to convert an OptimizedModule to a GraphModule directly if your primary goal is getting speedups, all you need to do is run an inference

However if you’re trying to modify the graph module in some way you can set a breakpoint here pytorch/compile_fx.py at master · pytorch/pytorch · GitHub and change the graphmodule that’s passed in

3

Thanks Mark!
But my goal is not just speedup inference but also trying to replace module like what I did in TorchFX, and according to official doc,

TorchDynamo generates FX Graphs from Python bytecode.

so I think maybe I can get a FX Graph from TorchDynamo, is there any API that allows me to do so?

4

This notebook might help Google Colab

5

Thanks! That’s very helpful!