Compute the disk size of a PyTorch model from its weights.

pytorch_model_disk_size.py

# 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