search

pyg-team / pytorch_geometric

Graph Neural Network Library for PyTorch
https://pyg.org
MIT License
21.17k stars 3.64k forks source link

Model summary functionality for PyG models #5727

Closed nowyouseemejoe closed 1 year ago

nowyouseemejoe commented 2 years ago

🚀 The feature, motivation and pitch

A model summary functionality to print a string summary of the network/model, like the tensorflow.Keras.Model.summary() method. A model summary would also help us debug the model and figure out the shapes of each layer's output. To be specific, summary can be implemented as a functional method, such that:

from torch_geometric.nn import summary
model = GCN(...)
summary(model, ...)

# Outputs (just an example):
_________________________________________________________________
Layer (type)                 Output Shape               #Param   
=================================================================
GCNConv1                   (-1, 16)              4253864   
_________________________________________________________________
GCNConv2                   (-1, 7)                  244       
=================================================================
Total params: 4,254,108
Trainable params: 4,232,220
Non-trainable params: 21,888
_________________________________________________________________

I saw there is a summary method for datasets, would there be a summary method for models? Thanks!

Alternatives

FYI:

Additional context

No response

rusty1s commented 2 years ago

Does there exist some external package that already supports this? Otherwise, I am happy to include this one into PyG.

nowyouseemejoe commented 2 years ago

Found an external package torchinfo that is not directly applicable to PyG-like inputs though. It would be amazing if PyG support this feature. Thanks!

rusty1s commented 2 years ago

I see. I guess this one fails since shapes are not static in PyG. In that case, I am supportive of the idea on integrating it directly into PyG. Let me know if you want to work on this together.

EdisonLeeeee commented 2 years ago

Sounds interesting. I can help with this as well :)

nowyouseemejoe commented 2 years ago

That would be great. Thank you @rusty1s @EdisonLeeeee

rusty1s commented 1 year ago

This is now available thanks to @EdisonLeeeee