New Tracing Mechanism

[mert-kurttutan]

Hi,

I want to have a discussion about the tracing mechanism used in torchinfo.
Current, it traces modules by registering hooks on modules.

One disadvantage of using only modules when tracing is that torch function that are not inside capture by modules, are not traced by this. For instance, addition using + operation for skip connection is not recorded, see #126 for torch.mul.

I think with the improved __torch_function__ and torch.Tensor subclassing features, we can trace better and also capture torch.functions (in addition to modules). (see here for intro)

One example of such application of __torch_function__ and torch.Tensor can be seen my project here.

There, this tracing mechanism is able to capture all the operations including arithmetic ops between tensors, modules, and tensor creation etc.

One downside is that it does not capture anything in jit traced modules, - I guess __torch_function__ does not work in traced modules. For traced modules, my strategy would be keep using hooks since apparently registering hooks still works on traced modules.

@TylerYep I used some parts of code in torchinfo project that I referenced at the end of README file. If you want add more citation or modification, you are welcome to do so!

[TylerYep]

Hey, thanks for opening up this discussion!

I think this is really interesting idea, and has the opportunity to refine the calculations of macs memory usage greatly. That being said, this would require a pretty foundational rewrite of the torchinfo library, and likely would not be as stable as the current implementation for some time. Would the layers shown by torchinfo change if we showed operator granularity? It likely won't be easy to show these operations as text compared to your torchview project.

If you are interested in adding this to torchinfo, I would recommend developing it behind a feature flag and having the output tests compare the old vs new output (I think only the memory usage statistics would change at first, the layers would stay the same). If we can prove that the new tracing system yields strictly better results than the old tracing system over all test cases, I would be happy to switch and release this as version 2.0.0.

[mert-kurttutan]

Yeah, I agree that it will certainly take some time. At the moment, I can slowly start looking into it.
I also encourage other volunteers to take a look at this issue (and my project for an example usage of __torch_function__).

Q: Would the layers shown by torchinfo change if we showed operator granularity?
A: I think it depends on the implementation. If we choose to include operator info (which I think we should), then those operators would be shown. Even if we don't include them in the representation, it will help calculating MACs and Memory Consumption of torch operators.

[johnmarktaylor91]

Not to self-promote, but I recently released a package, TorchLens , that uses a tracing mechanism for extracting metadata and activations from arbitrary neural networks, in addition to visualizing their structure. It’s currently slower and less optimized than torchinfo, but it can log metadata about any PyTorch operation, not just the modules.