前几天跑了一个大语言模型的代码，自家的电脑显卡本身本身是可以放下模型的参数和优化器参数的，但是训练的时候就报错，当时没有多想，就直接在云服务器上租了一个A800的显卡，不过今天突然想到了这个问题了。
看到了这么一个项目：
https://github.com/zhongzhengli13/MobileNetV3-for-leaf
其中的模型定义代码：
import torch
import torch.nn as nn
from torchvision.models import mobilenet_v3_small
from torchsummary import summary
class PlantDiseaseClassifier(nn.Module):
def __init__(self, num_classes=3):
super(PlantDiseaseClassifier, self).__init__()
self.base_model = mobilenet_v3_small(pretrained=False)
in_features = self.base_model.classifier[3].in_features
print(self.base_model.classifier) # 打印原始分类器结构
self.base_model.classifier[3] = nn.Linear(in_features, num_classes)
def forward(self, x):
return self.base_model(x)
if __name__ == "__main__":
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = PlantDiseaseClassifier(num_classes=3).to(device)
# ✅ 把 dummy_input 移动到相同 device 上
dummy_input = torch.randn(4, 3, 224, 224).to(device)
# 测试 forward
output = model(dummy_input)
print("\n✅ 输出形状：", output.shape) # 应该是 [4, 3]
# 模型结构 summary
summary(model, (3, 4000, 2672), device=str(device))
运行结果：
Sequential(
(0): Linear(in_features=576, out_features=1024, bias=True)
(1): Hardswish()
(2): Dropout(p=0.2, inplace=True)
(3): Linear(in_features=1024, out_features=1000, bias=True)
)
✅ 输出形状： torch.Size([4, 3])
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 16, 2000, 1336] 432
BatchNorm2d-2 [-1, 16, 2000, 1336] 32
Hardswish-3 [-1, 16, 2000, 1336] 0
Conv2d-4 [-1, 16, 1000, 668] 144
BatchNorm2d-5 [-1, 16, 1000, 668] 32
ReLU-6 [-1, 16, 1000, 668] 0
AdaptiveAvgPool2d-7 [-1, 16, 1, 1] 0
Conv2d-8 [-1, 8, 1, 1] 136
ReLU-9 [-1, 8, 1, 1] 0
Conv2d-10 [-1, 16, 1, 1] 144
Hardsigmoid-11 [-1, 16, 1, 1] 0
SqueezeExcitation-12 [-1, 16, 1000, 668] 0
Conv2d-13 [-1, 16, 1000, 668] 256
BatchNorm2d-14 [-1, 16, 1000, 668] 32
InvertedResidual-15 [-1, 16, 1000, 668] 0
Conv2d-16 [-1, 72, 1000, 668] 1,152
BatchNorm2d-17 [-1, 72, 1000, 668] 144
ReLU-18 [-1, 72, 1000, 668] 0
Conv2d-19 [-1, 72, 500, 334] 648
BatchNorm2d-20 [-1, 72, 500, 334] 144
ReLU-21 [-1, 72, 500, 334] 0
Conv2d-22 [-1, 24, 500, 334] 1,728
BatchNorm2d-23 [-1, 24, 500, 334] 48
InvertedResidual-24 [-1, 24, 500, 334] 0
Conv2d-25 [-1, 88, 500, 334] 2,112
BatchNorm2d-26 [-1, 88, 500, 334] 176
ReLU-27 [-1, 88, 500, 334] 0
Conv2d-28 [-1, 88, 500, 334] 792
BatchNorm2d-29 [-1, 88, 500, 334] 176
ReLU-30 [-1, 88, 500, 334] 0
Conv2d-31 [-1, 24, 500, 334] 2,112
BatchNorm2d-32 [-1, 24, 500, 334] 48
InvertedResidual-33 [-1, 24, 500, 334] 0
Conv2d-34 [-1, 96, 500, 334] 2,304
BatchNorm2d-35 [-1, 96, 500, 334] 192
Hardswish-36 [-1, 96, 500, 334] 0
Conv2d-37 [-1, 96, 250, 167] 2,400
BatchNorm2d-38 [-1, 96, 250, 167] 192
Hardswish-39 [-1, 96, 250, 167] 0
AdaptiveAvgPool2d-40 [-1, 96, 1, 1] 0
Conv2d-41 [-1, 24, 1, 1] 2,328
ReLU-42 [-1, 24, 1, 1] 0
Conv2d-43 [-1, 96, 1, 1] 2,400
Hardsigmoid-44 [-1, 96, 1, 1] 0
SqueezeExcitation-45 [-1, 96, 250, 167] 0
Conv2d-46 [-1, 40, 250, 167] 3,840
BatchNorm2d-47 [-1, 40, 250, 167] 80
InvertedResidual-48 [-1, 40, 250, 167] 0
Conv2d-49 [-1, 240, 250, 167] 9,600
BatchNorm2d-50 [-1, 240, 250, 167] 480
Hardswish-51 [-1, 240, 250, 167] 0
Conv2d-52 [-1, 240, 250, 167] 6,000
BatchNorm2d-53 [-1, 240, 250, 167] 480
Hardswish-54 [-1, 240, 250, 167] 0
AdaptiveAvgPool2d-55 [-1, 240, 1, 1] 0
Conv2d-56 [-1, 64, 1, 1] 15,424
ReLU-57 [-1, 64, 1, 1] 0
Conv2d-58 [-1, 240, 1, 1] 15,600
Hardsigmoid-59 [-1, 240, 1, 1] 0
SqueezeExcitation-60 [-1, 240, 250, 167] 0
Conv2d-61 [-1, 40, 250, 167] 9,600
BatchNorm2d-62 [-1, 40, 250, 167] 80
InvertedResidual-63 [-1, 40, 250, 167] 0
Conv2d-64 [-1, 240, 250, 167] 9,600
BatchNorm2d-65 [-1, 240, 250, 167] 480
Hardswish-66 [-1, 240, 250, 167] 0
Conv2d-67 [-1, 240, 250, 167] 6,000
BatchNorm2d-68 [-1, 240, 250, 167] 480
Hardswish-69 [-1, 240, 250, 167] 0
AdaptiveAvgPool2d-70 [-1, 240, 1, 1] 0
Conv2d-71 [-1, 64, 1, 1] 15,424
ReLU-72 [-1, 64, 1, 1] 0
Conv2d-73 [-1, 240, 1, 1] 15,600
Hardsigmoid-74 [-1, 240, 1, 1] 0
SqueezeExcitation-75 [-1, 240, 250, 167] 0
Conv2d-76 [-1, 40, 250, 167] 9,600
BatchNorm2d-77 [-1, 40, 250, 167] 80
InvertedResidual-78 [-1, 40, 250, 167] 0
Conv2d-79 [-1, 120, 250, 167] 4,800
BatchNorm2d-80 [-1, 120, 250, 167] 240
Hardswish-81 [-1, 120, 250, 167] 0
Conv2d-82 [-1, 120, 250, 167] 3,000
BatchNorm2d-83 [-1, 120, 250, 167] 240
Hardswish-84 [-1, 120, 250, 167] 0
AdaptiveAvgPool2d-85 [-1, 120, 1, 1] 0
Conv2d-86 [-1, 32, 1, 1] 3,872
ReLU-87 [-1, 32, 1, 1] 0
Conv2d-88 [-1, 120, 1, 1] 3,960
Hardsigmoid-89 [-1, 120, 1, 1] 0
SqueezeExcitation-90 [-1, 120, 250, 167] 0
Conv2d-91 [-1, 48, 250, 167] 5,760
BatchNorm2d-92 [-1, 48, 250, 167] 96
InvertedResidual-93 [-1, 48, 250, 167] 0
Conv2d-94 [-1, 144, 250, 167] 6,912
BatchNorm2d-95 [-1, 144, 250, 167] 288
Hardswish-96 [-1, 144, 250, 167] 0
Conv2d-97 [-1, 144, 250, 167] 3,600
BatchNorm2d-98 [-1, 144, 250, 167] 288
Hardswish-99 [-1, 144, 250, 167] 0
AdaptiveAvgPool2d-100 [-1, 144, 1, 1] 0
Conv2d-101 [-1, 40, 1, 1] 5,800
ReLU-102 [-1, 40, 1, 1] 0
Conv2d-103 [-1, 144, 1, 1] 5,904
Hardsigmoid-104 [-1, 144, 1, 1] 0
SqueezeExcitation-105 [-1, 144, 250, 167] 0
Conv2d-106 [-1, 48, 250, 167] 6,912
BatchNorm2d-107 [-1, 48, 250, 167] 96
InvertedResidual-108 [-1, 48, 250, 167] 0
Conv2d-109 [-1, 288, 250, 167] 13,824
BatchNorm2d-110 [-1, 288, 250, 167] 576
Hardswish-111 [-1, 288, 250, 167] 0
Conv2d-112 [-1, 288, 125, 84] 7,200
BatchNorm2d-113 [-1, 288, 125, 84] 576
Hardswish-114 [-1, 288, 125, 84] 0
AdaptiveAvgPool2d-115 [-1, 288, 1, 1] 0
Conv2d-116 [-1, 72, 1, 1] 20,808
ReLU-117 [-1, 72, 1, 1] 0
Conv2d-118 [-1, 288, 1, 1] 21,024
Hardsigmoid-119 [-1, 288, 1, 1] 0
SqueezeExcitation-120 [-1, 288, 125, 84] 0
Conv2d-121 [-1, 96, 125, 84] 27,648
BatchNorm2d-122 [-1, 96, 125, 84] 192
InvertedResidual-123 [-1, 96, 125, 84] 0
Conv2d-124 [-1, 576, 125, 84] 55,296
BatchNorm2d-125 [-1, 576, 125, 84] 1,152
Hardswish-126 [-1, 576, 125, 84] 0
Conv2d-127 [-1, 576, 125, 84] 14,400
BatchNorm2d-128 [-1, 576, 125, 84] 1,152
Hardswish-129 [-1, 576, 125, 84] 0
AdaptiveAvgPool2d-130 [-1, 576, 1, 1] 0
Conv2d-131 [-1, 144, 1, 1] 83,088
ReLU-132 [-1, 144, 1, 1] 0
Conv2d-133 [-1, 576, 1, 1] 83,520
Hardsigmoid-134 [-1, 576, 1, 1] 0
SqueezeExcitation-135 [-1, 576, 125, 84] 0
Conv2d-136 [-1, 96, 125, 84] 55,296
BatchNorm2d-137 [-1, 96, 125, 84] 192
InvertedResidual-138 [-1, 96, 125, 84] 0
Conv2d-139 [-1, 576, 125, 84] 55,296
BatchNorm2d-140 [-1, 576, 125, 84] 1,152
Hardswish-141 [-1, 576, 125, 84] 0
Conv2d-142 [-1, 576, 125, 84] 14,400
BatchNorm2d-143 [-1, 576, 125, 84] 1,152
Hardswish-144 [-1, 576, 125, 84] 0
AdaptiveAvgPool2d-145 [-1, 576, 1, 1] 0
Conv2d-146 [-1, 144, 1, 1] 83,088
ReLU-147 [-1, 144, 1, 1] 0
Conv2d-148 [-1, 576, 1, 1] 83,520
Hardsigmoid-149 [-1, 576, 1, 1] 0
SqueezeExcitation-150 [-1, 576, 125, 84] 0
Conv2d-151 [-1, 96, 125, 84] 55,296
BatchNorm2d-152 [-1, 96, 125, 84] 192
InvertedResidual-153 [-1, 96, 125, 84] 0
Conv2d-154 [-1, 576, 125, 84] 55,296
BatchNorm2d-155 [-1, 576, 125, 84] 1,152
Hardswish-156 [-1, 576, 125, 84] 0
AdaptiveAvgPool2d-157 [-1, 576, 1, 1] 0
Linear-158 [-1, 1024] 590,848
Hardswish-159 [-1, 1024] 0
Dropout-160 [-1, 1024] 0
Linear-161 [-1, 3] 3,075
MobileNetV3-162 [-1, 3] 0
================================================================
Total params: 1,520,931
Trainable params: 1,520,931
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 122.31
Forward/backward pass size (MB): 7357.33
Params size (MB): 5.80
Estimated Total Size (MB): 7485.44
----------------------------------------------------------------
上面的信息中重点在于下面的内容：
可以看到这么一个参数大小只有5.8MB显存大小的CNN小模型，加上优化器的参数也就12MB大小不到，为什么在测试中显存总共占了7485.44MB，这个是7.4GB的大小，这个和大模型训练时候的问题是一样的，为什么在最终训练时候显存占用远远高于模型的大小。
问下豆包大模型，给出AI生成的答案：
这个答案的可信度还是比较高的，为此我这里将图片大小调整为：
(3, 224, 224)
也就是说上面的代码中将图片输入部分代码替换为：
# 模型结构 summary
summary(model, (3, 224, 224), device=str(device))
再次运行，给出结果：
可以看到，这次的显存总共占用为40.98MB大小，这也一定程度上解释了大模型训练过程中显存不是比模型大2倍那么简单的问题，由于现在的大模型算法很多都是长文本加Reinforcement Learning微调，这样就导致输入的数据会是一个非常大的size，这样就出现了本文所提到的问题，这也是为什么一个0.5B大小的大模型需要使用一个A800显卡来训练，80GB的显存被占掉70GB多，这是同样的问题，那就是训练过程中如果一次训练（forward/backforward）过程中输入数据过大会导致总显存的占用会远远大于模型参数大小。
本文中最初的大size的图片输入所占用的显存结构如下：
输入(122MB) + 特征/梯度(7357MB) + 参数(5.8MB) ≈ 7485MB
可以看到上面的显存中绝大部分的显存都是因为输入数据过大导致前后向计算中中间数据所导致的，也就是说计算过程中的显存峰值远远大于模型的显存占用大小。
本博客是博主个人学习时的一些记录，不保证是为原创，个别文章加入了转载的源地址，还有个别文章是汇总网上多份资料所成，在这之中也必有疏漏未加标注处，如有侵权请与博主联系。 如果未特殊标注则为原创，遵循 CC 4.0 BY-SA 版权协议。
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。