如何找到深度学习模型的大小？

Question

我正在使用同一模型的不同量化实现，主要区别在于权重、偏差和激活的精度。所以我想知道如何找到 32 位浮点数的模型大小（以 MB 为单位）和 int8 模型的大小之间的差异。我已将模型保存为 .PTH 格式。

Answer 1

您可以计算参数和缓冲区的数量。 然后尝试将它们与元素大小相乘，您将得到所有参数的大小。

model = models.resnet18()
param_size = 0
buffer_size = 0
for param in model.parameters():
    param_size += param.nelement() * param.element_size()

for buffer in model.buffers():
    buffer_size += buffer.nelement() * buffer.element_size()

size_all_mb = (param_size + buffer_size) / 1024**2
print('Size: {:.3f} MB'.format(size_all_mb))

它会打印：

大小：361.209 MB

Answer 2

“要计算模型大小（以字节为单位），请将参数数量乘以所选精度的大小（以字节为单位）。例如，如果我们使用 BLOOM-176B 模型的 bfloat16 版本，则有 176*10**9 x 2 字节 = 352GB！”

关于 HF 的这篇博客值得一读：https://huggingface.co/blog/hf-bitsandbytes-integration

Answer 3

编写了一段小代码来根据模型的参数数量和模型的

dtype

来计算模型的大小
目前支持

fp32

、

fp16

和

int8

def cal_size(num_params,dtype):
    if dtype == "float32":
        return (num_params/10**6) * 4
    elif dtype == "float16":
        return (num_params/10**6) * 2
    elif dtype == "int8":
        return (num_params/10**6) * 1
    else:
        return -1

if __name__ == "__main__":
    import torchvision.models as models
    model = models.mobilenet_v2()
    #mobilenetv2 of width multiplier 1 has 3.4M params 
    total_params = sum(p.numel() for p in model.parameters())
    model_size = cal_size(total_params,"float32")
    if model_size != -1:
        print("Size of model is :{:.2f} MB".format(model_size))
    else:
        print("Incorrect dtype")