Mapping between the node in onnx to the pytorch module

Hi~ Is there any way to map the nodes in the onnx.graph(exported by torch.onnx.export) to the modules in the pytorch module(for example, conv). The reason I do this is that I need some additional customized attributes in the module to optimize the execution of onnx, but this information is not included in onnx (this information is not used in inference).

1 Like