search

delve-team / delve

PyTorch model training and layer saturation monitor
https://delve-docs.readthedocs.io
MIT License
79 stars 13 forks source link

Improvement: Can SaturationTracker be added to only track layers in the specified list #66

Closed liujuncn closed 1 year ago

liujuncn commented 1 year ago

Too many layers can cause GPU OOM when tracking large models. So is it possible to track only some of all layers.

For example,

tracker = SaturationTracker( modules=some_model, layers = [layer_1, layer_4, ..., layer_last] )

MLRichter commented 1 year ago

Hey, thanks for your feedback. I think it should be fairly straightforward to create a filter which results in specific hooks not being attached undesired layers. I will look into it.

MLRichter commented 1 year ago

Okay, got it, checkout this branch https://github.com/delve-team/delve/tree/feature/layer_filter SaturationTracker now has an additional argument layer_filter, which accepts a user-defined function for manipulating the dictionary of tracked layers. The tracked layers are held in a dictionary mapping their names (strings) to the nn.Module-objects.

For installation, please clone the branch, navigate into the repositories folder and pip install -e . to install this new version. If this solves your issue, I will include it in the next release.

liujuncn commented 1 year ago

Yes, the filter works.

Wait for your next release.

MLRichter commented 1 year ago

Release 0.1.50 is out now, containing the layer filters, you requested https://pypi.org/project/delve/

thank you for your constructive feedback.