深度学习代码块之计算模型参数量和显存大小

本文链接：https://blog.csdn.net/weixin_48235955/article/details/136078535

🙋大家好！我是毛毛张!
🌈个人首页： 神马都会亿点点的毛毛张

今天毛毛张介绍一个关于如何计算模型参数量和显存大小的深度学习代码块

文章目录

0 引言
1 构建模型
方法1：函数封装
方法2：迭代
方法3：使用第三方库
- 1. torchstat库
- 2. thop库
总结
参考文献

0 引言

我们听说的llama2-7B大模型，是指该模型的参数量有70亿，那么如何进行计算呢，以及如何计算70亿的参数量加载到显存中需要多少内存。通常模型中每个参数默认是按单精度浮点数（FP32）进行存储，该格式4个字节(4Bytes)，同时有换算关系：1M = 1024 KB = 1024*1024 B，所以只要我们能够知道模型的参数量，我们就能够计算出模型所占用的显存大小。下面毛毛张将主要介绍如何进行计算模型的参数量大小，毛毛张首先会构建一个GoogleNet模型，然后围绕计算该网络模型的参数量。

1 构建模型

import torch.optim
import torch
from torch import nn
import torch.nn.functional as F

class InceptionA(nn.Module):
    def __init__(self, in_channels):
        super(InceptionA, self).__init__()
        self.branch1X1 = nn.Conv2d(in_channels, 16, kernel_size=1)

        # 设置padding保证各个分支输出的高度和宽度保持不变
        self.branch5X5_1 = nn.Conv2d(in_channels, 16, kernel_size=1)
        self.branch5X5_2 = nn.Conv2d(16, 24, kernel_size=5, padding=2)

        self.branch3X3_1 = nn.Conv2d(in_channels, 16, kernel_size=1)
        self.branch3X3_2 = nn.Conv2d(16, 24, kernel_size=3, padding=1)
        self.branch3X3_3 = nn.Conv2d(24, 24, kernel_size=3, padding=1)

        self.branch_pool = nn.Conv2d(in_channels, 24, kernel_size=1)

    def forward(self, x):
        branch1X1 = self.branch1X1(x)

        branch5X5 = self.branch5X5_1(x)
        branch5X5 = self.branch5X5_2(branch5X5)

        branch3X3 = self.branch3X3_1(x)
        branch3X3 = self.branch3X3_2(branch3X3)
        branch3X3 = self.branch3X3_3(branch3X3)

        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1)
        branch_pool = self.branch_pool(branch_pool)

        outputs = [branch1X1, branch5X5, branch3X3, branch_pool]
        # （b, c, w, h），dim=1——以第一个维度channel来拼接
        return torch.cat(outputs, dim=1)


# 定义模型
class GoogleNet(nn.Module):
    def __init__(self):
        super(GoogleNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        # 88 = 24*3 + 16
        self.conv2 = nn.Conv2d(88, 20, kernel_size=5)

        self.incep1 = InceptionA(in_channels=10)
        self.incep2 = InceptionA(in_channels=20)

        self.mp = nn.MaxPool2d(2)
        # 确定输出张量的尺寸
        # 在定义时先不定义fc层，随便选取一个输入，经过模型后查看其尺寸
        # 在init函数中把fc层去掉，forward函数中把最后两行去掉，确定输出的尺寸后再定义Lear层的大小
        self.fc = nn.Linear(1408, 10)

    def forward(self, x):
        in_size = x.size(0)
        # (C,W,H)=(1,28,28)
        x = F.relu(self.mp(self.conv1(x))) 
        # (C,W,H)=(10,12,12)
        x = self.incep1(x) 
        # (C,W,H)=(88,12,12)
        x = F.relu(self.mp(self.conv2(x))) 
        # (C,W,H)=(20,4,4)
        x = self.incep2(x)  
        # (C,W,H)=(88,4,4)
        x = x.view(in_size, -1)  # 1408=4*4*88
        x = self.fc(x)
        return x

方法1：函数封装

def get_memory_size(model, bytes_per_param=4):
    """
    该函数用于计算模型占用的显存大小，根据每个参数的字节数进行计算。

    :param model: 实例化的模型
    :param bytes_per_param: 每个参数占用的字节数，默认是按全精度，FP32精度，4个字节
    :return: None
    """
    # 计算模型中所有参数的总数量
    total_params = sum(p.numel() for p in model.parameters())
    total_params_memory_size = total_params * bytes_per_param
    # 这段代码的作用是统计模型中所有需要梯度（可训练）的参数的数量
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    trainable_params_memory_size = trainable_params * bytes_per_param

    print(f'Total parameters: {total_params}')
    print(f'total_params_memory_size: {total_params_memory_size / (1024 * 1024):.2f} MB')
    print(f'Total trainable parameters : {trainable_params}')
    print(f'trainable_params_memory_size: {trainable_params_memory_size / (1024 * 1024):.2f} MB')


model = GoogleNet()
get_memory_size(model)
# 输出：
"""
Total parameters: 97298
total_params_memory_size: 0.37 MB
Total trainable parameters : 97298
trainable_params_memory_size: 0.37 MB
"""

方法2：迭代

model = GoogleNet().to(device)
# 获取模型中的所有参数，并存储在列表 params 中
params = list(model.parameters())
# params 中的每个元素都是模型的一个参数张量，包含了参数的值以及其他相关信息，比如参数的形状（shape）等
# 可以通过调试来来查看params的元素的具体属性，有助于理解下面代码
# 初始化变量 k，用于累加所有参数的元素数量
k = 0

# 对参数列表中的每个参数张量进行迭代
for i in params:
    # 初始化变量 l，用于存储当前参数张量的元素数量
    l = 1
    # 对当前参数张量的形状中的每个维度进行迭代
    for j in i.size():
        # 计算当前维度上的元素数量，并将其累乘到变量 l 中
        l *= j
    # 将当前参数张量的元素数量累加到变量 k 中
    k = k + l

# 输出模型中所有参数的总数量
print("总参数数量和：" + str(k))
# 输出：
"""
总参数数量和：97298
"""

方法3：使用第三方库

1. torchstat库

from torchstat import stat
model = GoogleNet()
stat(model, (1, 28, 28))
"""
===================================================================================================
Total params: 97,298
--------------------------------------------------------------------------------------------------
Total memory: 0.12MB
Total MAdd: 12.04MMAdd
Total Flops: 6.05MFlops
Total MemR+W: 624.42KB
"""

2. thop库

from thop import profile
model = GoogleNet()
flops, params = profile(model, inputs=(torch.randn(1, 1, 28, 28), ))
print(f"FLOPs: {flops}, Parameters: {params}")
# 输出：
"""
FLOPs: 6019200.0, Parameters: 97298.0
"""