Pytorch-lightning: Torch-Summary Integration to replace ModelSummary

[TylerYep]

🚀 Feature, Motivation, Pitch

Hello, I am the current maintainer of torch-summary, which is a rewrite of yet another torchsummary library. Currently, ModelSummary is a builtin PyTorch-Lightning feature, but it is limited in scope and support, especially in more complex models. The proposal here, seconded by @carmocca , is to deprecate/remove the current ModelSummary implementation in pytorch-lightning and instead move to torch-summary as a more sustainable and long-term solution to the problem of visualizing model details.

Implementations

There are several ways to implement this idea, all of which I am open to:

• Using torch-summary as a required dependency for pytorch-lightning by adding it to requirements.txt and integrating it into applicable models. Lightning would use the current public API for torch-summary.

• Bring the torch-summary code in-house as an upgrade to the existing code as a one time upgrade.

• Adopt torch-summary (or a copy of the repository) into the PytorchLightning organization as a Lightning-specific repository that can be developed independently of the original project. (in other words, surfacing an API that only works for Lightning models). I hope that this would add additional contributors that will help the project grow to fit more diverse kinds of models.

Additional context

As the maintainer, I will say that this project is far from perfect, and there are many kinds of models that torch-summary does not yet fully support. e.g. functional layers, lists or dicts of tensors, other types of CUDA, etc. The project is currently only maintained by me, which is sustainable for the current user base, but not sustainable if Lightning users' needs eventually outpace the project.

However, I do think that separating the ModelSummary functionality from Lightning is important, and I think that expanding this feature is something that a lot of users would enjoy.

(Continuing discussion from #4521)

[awaelchli]

Hi, I allowed myself to make an edit to the description because I do not actually second the removal of the current summary.

The current summary code is lightweight, simple and has little Lightning specific code. I highly doubt that it will be possible to add a hard dependency on an external repo just for a summary text. So option A is probably not feasible. Note: at the beginning, we had pandas as a strict dependency to format the table, but had to remove it because it was too heavy.

_Disclaimer: I am probably biased because I wrote some of the essentials of the model summary code._

I am generally open to B and C.

I see one more option: keep current summary for simple printing and dependency free operation, and if the user wants torch summary, provide it as a callback that can be easily plugged in given user has installed the dependency.

[carmocca]

which is a rewrite of yet another torchsummary library

How much is it a whole rewrite and how much just maintaining previous code? the original repository has 2.7k stars (+114 in your fork) so there is definitely interest in this feature.

I highly doubt that it will be possible to add a hard dependency on an external repo just for a summary text. So option A is probably not feasible. Note: at the beginning, we had pandas as a strict dependency to format the table, but had to remove it because it was too heavy.

I havent dived into the torch-summary code but I expect it to be quite light. I don't think comparing it to pandas is fair.

The reason I proposed this is because #4521 adds one of its features so, I thought, why not join efforts and use the code that is already written?

The current summary code is lightweight, simple

It might be right now, but it might not be in the future. I think having it separate from Lightning is a great idea (outside or inside the Lightning organization)

So if A is not an option, IMO C would be better than B

I see one more option: keep current summary for simple printing and dependency free operation, and if the user wants torch summary, provide it as a callback that can be easily plugged in given user has installed the dependency.

What if -- with time -- users start making PRs to current summary and then you find you converged into the extra dependency?

[TylerYep]

Ah, sorry about that @awaelchli , I misread your comment on the linked issue.

How much is it a whole rewrite and how much just maintaining previous code?

Torch-summary, in my opinion, is a complete rewrite - while the output is similar, the underlying implementation is unrecognizable compared to the original.

I haven't dived into the torch-summary code but I expect it to be quite light. I don't think comparing it to pandas is fair.

I would agree that torch-summary is very lightweight - there are currently only 4 source files in total, averaging around 200 lines of code each.

I think option C makes more sense given Pytorch-Lightning's needs. It gives the functionality + API room to grow separately from the version meant for regular PyTorch users. It is also a good idea to develop it separately from the main Lightning repo because it is not guaranteed to work on all Pytorch-Lightning models, and thus may require a different release schedule.

[tchaton]

Hey @williamFalcon, could you have a look at this one ?