Skip to content

Instantly share code, notes, and snippets.

@lparolari
Last active July 11, 2023 12:49
Show Gist options
  • Save lparolari/f2a74f8a19c38c6d70afeea5b91b8b7d to your computer and use it in GitHub Desktop.
Save lparolari/f2a74f8a19c38c6d70afeea5b91b8b7d to your computer and use it in GitHub Desktop.
Compute the disk size of a PyTorch model from its weights.
# Reference: https://discuss.pytorch.org/t/why-are-saved-models-so-large/166324/2
for name, param in model.named_parameters():
size = int((param.nelement() * param.element_size()) / (1024 ** 2)) # MB = 1024^2 bytes
print(name, param.nelement(), param.element_size(), f"{size}MB")
# EXAMPLE OUTPUT:
# concept_branch.we.embedding.weight 120005700 4 457MB
# visual_branch.we.embedding.weight 120005700 4 457MB
# visual_branch.fc.weight 78300 4 0MB
# visual_branch.fc.bias 300 4 0MB
# textual_branch.lstm.weight_ih_l0 360000 4 1MB
# textual_branch.lstm.weight_hh_l0 360000 4 1MB
# textual_branch.lstm.bias_ih_l0 1200 4 0MB
# textual_branch.lstm.bias_hh_l0 1200 4 0MB
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment