How to write jit model to a file in text format including backward?

jit
1

Hi,
I’m looking for a way to write the graph (preferably torch.jit.script module) forward and backward into a text format file.
A simple way to write the forward:
module = torch.jit.script(model) # JIT for model
print(module.code)
But, is there something similar to write the backward?
The reason I’m looking for that is: I want to evaluate performance of the forward and backward for a theoretical HW.
I tried to use pytorch profile, but it does not give all the information. Does not give the graph connectivity and tensor sizes
ONNX export support only forward.

Any other idea?

2

One thing that you could try is using the torchviz library to visualize what the backward pass looks like!

torchviz allows you to draw the computation graph of your network. This might be able to help you!