search

frgfm / torch-scan

Seamless analysis of your PyTorch models (RAM usage, FLOPs, MACs, receptive field, etc.)
https://frgfm.github.io/torch-scan/latest
Apache License 2.0
208 stars 22 forks source link

Extending support to custom modules #41

Open vkothapally opened 3 years ago

vkothapally commented 3 years ago

🚀 Compatibility with Complex Modules

Thank you for the great repo. Can you please let me know if you would be willing to extend this to be compatible with complex convolution modules too? (For instance, you can look at cplxmodule or speechBrain). This will really help is comparing various models from architectural details to memory usage standpoint.

frgfm commented 3 years ago

Hi @vkothapally,

Thanks for asking! There are two answers I can give you:

  1. As of now, this library is designed to detect inherited modules (from PyTorch parent classes) and leverages this to cover most of model's components. Using this design, it wouldn't be possible to add support for custom modules apart from manually adding dependencies to all libraries that define custom modules.

  2. That being said, I am considering completely changing the library for v0.2.0 by using ONNX export. It would limit support of ONNX exportable modules but I would be able to produce the same information without manually writing the computation of FLOPS, etc.

So in the short term, apart from if you have a good PR in mind, I'm afraid I won't add support to this before v0.2.0 :/

ratom commented 1 year ago

I tried to calculate the receptive fields for VGG, and ResNet. It works for VGG, but not for resnet. Can you help me.

frgfm commented 1 year ago

Hey @ratom,

As of now, the receptive field feature works on "highway nets" meaning that skip connections cannot be handled appropriately! I'll try to see how I can address this :+1: