Skip to content

Instantly share code, notes, and snippets.

@KeAWang

KeAWang/count_torch_params.py

Last active November 14, 2023 18:35
Show Gist options
  • Save KeAWang/020d51c67ddbac1728f0e9a283398aef to your computer and use it in GitHub Desktop.
Save KeAWang/020d51c67ddbac1728f0e9a283398aef to your computer and use it in GitHub Desktop.
Download ZIP
Count number of pytorch parameters
Raw
count_torch_params.py
import torch
def count_params(model: torch.nn.Module):
"""count number trainable parameters in a pytorch model"""
total_params = sum(torch.numel(x) for x in model.parameters())
return total_params
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment