Skip to content

Instantly share code, notes, and snippets.

@ihoromi4

ihoromi4/get_n_params.py

Created February 5, 2020 11:45
Show Gist options
  • Save ihoromi4/aa16085532358f9fc7937941526d827c to your computer and use it in GitHub Desktop.
Save ihoromi4/aa16085532358f9fc7937941526d827c to your computer and use it in GitHub Desktop.
Download ZIP
pytorch - get number of a model parameters
Raw
get_n_params.py
import functools
import operator
from torch import nn
def get_n_params(model: nn.Module) -> int:
return sum((functools.reduce(operator.mul, p.size()) for p in model.parameters()))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment