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
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?