Skip to content

Conversation

martijnvanbeers
Copy link

I wanted to be able to do things with the per-module flops data, so I wrote a little function to store them. To use it, I had to add a patch to be able to run my function through model.apply() before calculate_flops() removed all the data. Hopefully you find this useful too.

(I also cleaned up a bit and added the pyproject.toml so I could install from my repo)

Martijn van Beers added 3 commits December 5, 2024 15:39
copied from the tarball on pypi so I can (dev-)install my fork
Remove the compiled python files, and add __pycache__ to the .gitignore
so they won't come back
Add a parameter to be able to run extra functions through model.apply().
They are run right before the print function.

This lets you do things like store the FLOPS for each module rather than
just printing them, as the new example script shows.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant