Skip to content

Instantly share code, notes, and snippets.

@nmakes

nmakes/pytorch_count_parameters.py

Last active November 2, 2021 03:26
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save nmakes/52037e3c30c145bf1f1e15e0bb32badc to your computer and use it in GitHub Desktop.
Save nmakes/52037e3c30c145bf1f1e15e0bb32badc to your computer and use it in GitHub Desktop.
Download ZIP
Count number of parameters in a pytorch model
Raw
pytorch_count_parameters.py
'''
Count the number of parameters in a pytorch model.
Written for pytorch (cuda compatible).
Written by: Naveen Venkat (naveenvenkat.com)
'''
def count_params(model):
tot = 0
for p in model.parameters():
S = p.size(); n = 1
for x in S:
n *= x
tot += n
return tot
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment