
1. 一个让很多人困惑的现象模型文件明明只有几十兆加载到内存里却占了几百兆甚至上吉字节训练的时候更是直接把显存吃满。这个问题我在刚接触卷积神经网络的时候也踩过当时第一反应是“是不是框架有内存泄漏”排查了半天才发现问题根本不在框架而在于我对卷积运算的内存开销理解得太粗糙。这篇文章想做的事情很具体把卷积操作到底吃多少内存这件事从三个维度算清楚。参数量、MACs乘加运算次数、中间特征图占用这三笔账各自独立又相互关联。搞明白它们之间的关系你就能准确判断一个模型在运行时到底需要多少内存也能在模型设计和部署时做出更合理的取舍。适合读这篇内容的人包括正在学卷积神经网络但被内存问题困扰的开发者、需要在边缘设备上部署模型的工程师、以及想搞清楚“模型文件大小”和“运行时内存”为什么差距这么大的任何人。我会尽量用具体的数字和计算过程来说明而不是停留在概念层面。2. 先把三笔账的定义理清楚2.1 参数量模型文件大小的直接来源参数量就是模型中所有需要学习的权重和偏置的总数。对于一个标准的二维卷积层参数量的计算公式是参数量 卷积核高 × 卷积核宽 × 输入通道数 × 输出通道数 输出通道数偏置举个例子一个 3×3 卷积核输入通道 256输出通道 512那么参数量就是 3×3×256×512 512 1,179,136大约 118 万。如果用 FP32 存储每个参数占 4 字节这一层就是约 4.5 MB。这里有个容易被忽略的点参数量和模型文件大小基本对应。你保存模型时存的就是这些参数。所以如果模型文件是 50 MB那参数量大概就是 1300 万左右FP32。这个数字是静态的跟输入尺寸无关跟batch size也无关。2.2 MACs计算量的衡量标准MACs 全称是 Multiply-Accumulate operations即乘加运算次数。每一次“乘法加法”算作一次 MAC。卷积层的 MACs 计算公式MACs 卷积核高 × 卷积核宽 × 输入通道数 × 输出通道数 × 输出特征图高 × 输出特征图宽注意这里比参数量多乘了输出特征图的尺寸。也就是说MACs 和输入尺寸直接相关。同样一个卷积层输入从 224×224 变成 448×448MACs 会变成原来的 4 倍但参数量不变。这就是为什么模型文件小不代表计算量小。一个全卷积网络可能只有几兆参数但处理大分辨率图像时计算量可以非常惊人。2.3 中间特征图运行时内存的真正大头中间特征图是每一层卷积输出的激活值。它们在前向传播时需要保存在反向传播时还需要用来计算梯度。特征图的内存占用公式单层特征图内存 batch_size × 输出通道数 × 输出特征图高 × 输出特征图宽 × 每元素字节数用 FP32 的话每元素 4 字节。假设 batch_size32输出通道 512特征图 56×56那么这一层就是 32×512×56×56×4 205,520,896 字节约 196 MB。仅仅一层就接近 200 MB而这一层的参数量可能只有几 MB。三笔账的关系可以这样理解参数量决定模型文件大小MACs 决定计算耗时中间特征图决定运行时内存峰值。三者中中间特征图往往是内存占用的主要来源也是最容易被忽视的部分。3. 为什么模型文件小但运行时吃内存3.1 文件存的是参数内存里装的是全部模型文件里只有参数但运行时内存里装的东西远不止参数。具体来说包括模型参数和文件大小对应FP32 下每个参数 4 字节。中间激活值每一层的输出特征图前向传播时必须保留。梯度训练时每个参数对应一个梯度内存翻倍。优化器状态如果用 Adam每个参数还要存一阶矩和二阶矩再翻两倍。工作缓冲区卷积算法如 im2col、Winograd需要的临时空间。训练时一个 FP32 模型的总内存需求大约是参数量的 4 倍左右参数梯度优化器两个状态再加上激活值。推理时不需要梯度和优化器状态但激活值仍然存在。3.2 一个具体的计算例子拿一个简化的例子来算。假设有一个卷积层输入224×224×3卷积核3×3输出通道 64输出224×224×64same paddingbatch_size 16参数量3×3×3×64 64 1,792FP32 下约 7 KB。MACs3×3×3×64×224×224 86,704,128约 8670 万次。输出特征图内存16×64×224×224×4 205,520,896 字节约 196 MB。可以看到参数量只有 7 KB但输出特征图占了 196 MB。这就是“模型文件很小运行为什么还吃内存”的直接答案。3.3 实际网络中的累积效应单个层已经这么夸张整个网络累积起来更明显。以 VGG16 为例参数量约 1.38 亿FP32 下模型文件约 528 MB。但推理时如果输入 224×224第一层卷积输出就是 224×224×64batch_size1 时约 12.8 MB。随着层数加深通道数增加特征图尺寸减小但通道数增长更快。VGG16 的 conv4_3 层输出是 56×56×512单张图就是 56×56×512×4 6.4 MB。如果 batch_size16就是 102 MB。这还只是一层。整个网络所有层的激活值加起来峰值内存可能达到几百 MB甚至上吉字节。注意很多人只关注模型文件大小忽略了激活值的内存占用。在部署时激活值往往是限制 batch_size 和输入分辨率的关键因素。4. 卷积算法本身的内存开销4.1 im2col 的空间换时间策略直接做卷积在内存访问上效率很低因为要频繁跳跃读取。im2col 的思路是把输入特征图按照卷积核的感受野展开成一个大矩阵然后调用矩阵乘法。这样做的好处是矩阵乘法有高度优化的库如 BLAS计算速度快。但代价是内存。一个 3×3 卷积im2col 后矩阵的行数是输出特征图的位置数列数是 3×3×输入通道。对于 224×224×3 的输入输出 224×224×64im2col 矩阵大小是 (224×224) × (3×3×3) 50176 × 27约 135 万个元素FP32 下 5.4 MB。这还只是一层。如果输入通道是 256卷积核 3×3那么列数就是 2304矩阵大小变成 50176 × 2304约 1.16 亿个元素FP32 下 462 MB。这个临时矩阵就是 im2col 的内存开销。4.2 Winograd 的权衡Winograd 算法通过变换减少乘法次数但需要额外的变换矩阵和中间结果存储。对于 3×3 卷积Winograd F(2×2, 3×3) 可以把乘法次数减少到原来的 4/9 左右。但变换过程需要存储输入变换结果、权重变换结果和输出变换结果。这些中间结果的大小取决于 tile 的大小和通道数。在通道数很大时Winograd 的中间存储可能比 im2col 还大。所以很多框架会根据卷积参数自动选择算法在内存和速度之间做权衡。4.3 不同算法的内存对比算法计算效率额外内存适用场景直接卷积低很小小卷积核、小通道数im2col GEMM高大大通道数、通用Winograd很高中等3×3 卷积、通道数适中FFT中等大大卷积核实际框架如 PyTorch、TensorFlow 会根据卷积参数和硬件自动选择算法。但你可以通过设置环境变量或框架参数来限制算法选择从而控制内存。5. 训练和推理的内存差异5.1 推理时的内存构成推理时不需要保存梯度也不需要优化器状态。内存主要包括模型参数当前层的输入和输出特征图工作缓冲区如果框架做了内存复用比如 PyTorch 的 caching allocator不同层的特征图可以共享同一块内存。但峰值内存仍然由最大的那层决定。5.2 训练时的内存放大训练时内存需求大幅增加梯度每个参数一个梯度和参数同大小。优化器状态SGD 不需要额外状态Adam 需要一阶矩和二阶矩每个参数两个额外值。激活值反向传播需要用到前向传播的中间结果所以所有层的激活值都要保存。以一个 1000 万参数的模型为例项目FP32 大小参数40 MB梯度40 MBAdam 一阶矩40 MBAdam 二阶矩40 MB激活值假设200 MB合计360 MB如果 batch_size 翻倍激活值翻倍总内存变成 560 MB。这就是为什么增大 batch_size 会很快耗尽显存。5.3 混合精度的影响用 FP16 代替 FP32 可以把参数、梯度、激活值的内存都减半。但优化器状态通常还是 FP32所以总内存不是简单减半。实际测试中混合精度训练通常能节省 30% 到 50% 的内存。提示混合精度训练时loss scaling 是必须的否则梯度会下溢。PyTorch 的 amp 模块可以自动处理。6. 实际排查内存问题的思路6.1 先算理论值再对比实际值遇到内存问题时先按上面的公式算一下理论峰值。如果实际内存远大于理论值可能是以下原因框架没有做内存复用算法选择导致额外内存数据加载器占用过多内存内存碎片6.2 用工具定位PyTorch 可以用torch.cuda.memory_summary()查看显存分配情况。TensorFlow 可以用tf.config.experimental.get_memory_info()。对于 CPU 内存可以用memory_profiler或tracemalloc。6.3 常见优化手段减小 batch_size最直接但可能影响训练稳定性。使用梯度累积用小 batch 模拟大 batch。梯度检查点用计算换内存不保存所有激活值反向传播时重新计算。混合精度FP16 减半内存。模型剪枝减少参数和计算量。知识蒸馏用大模型教小模型。6.4 一个实际案例我之前部署一个分割网络到边缘设备模型文件只有 20 MB但推理时内存占用超过 500 MB。排查后发现输入分辨率是 1024×1024第一层输出就是 1024×1024×64FP32 下 256 MB。框架默认用了 im2col额外分配了 100 MB 的临时矩阵。没有做内存复用每一层都新分配内存。解决方案把输入降到 512×512内存降到 120 MB启用内存复用再降到 80 MB最后用 FP16 推理降到 45 MB。模型精度只下降了 0.5%。7. 一些容易踩的坑和实操建议7.1 不要只看模型文件大小模型文件大小只反映参数量不反映运行时内存。一个全卷积网络可能只有几兆参数但处理大图时激活值可以占几个吉字节。评估模型内存需求时一定要算激活值。7.2 注意框架的默认行为不同框架的内存管理策略不同。PyTorch 的 caching allocator 会缓存已分配的内存导致nvidia-smi显示的显存占用比实际使用高。这是正常现象不是内存泄漏。7.3 输入分辨率的影响是平方级的输入从 224 变成 448特征图面积变成 4 倍激活值内存也变成 4 倍。在边缘设备上降低输入分辨率是最有效的内存优化手段。7.4 通道数的影响是线性的输出通道数翻倍激活值内存翻倍参数量也翻倍。但通道数对精度的影响通常比分辨率更敏感所以优先降分辨率而不是降通道数。7.5 用分组卷积和深度可分离卷积深度可分离卷积把标准卷积分解为深度卷积和逐点卷积参数量和计算量都大幅减少。MobileNet 系列就是基于这个思路。但深度可分离卷积的激活值内存不一定减少因为逐点卷积的输出通道数可能很大。7.6 注意数据加载器的内存很多人忽略了 DataLoader 的内存占用。如果num_workers设置过大每个 worker 都会复制一份数据内存成倍增加。通常num_workers设置为 CPU 核心数即可不要盲目调大。7.7 内存碎片问题长时间训练时频繁分配和释放不同大小的内存块会导致碎片。PyTorch 的 caching allocator 可以缓解这个问题但在某些情况下仍然需要手动调用torch.cuda.empty_cache()。不过这个操作会降低性能不要频繁调用。8. 从三笔账到实际决策回到最初的问题模型文件很小运行为什么还吃内存因为文件里只有参数而运行时内存里还有激活值、梯度、优化器状态和工作缓冲区。参数量、MACs、中间特征图这三笔账各自独立但共同决定了模型的实际资源需求。在实际项目中我通常会先算这三笔账然后决定如果参数量大但激活值小考虑模型剪枝或量化。如果激活值大但参数量小考虑降低输入分辨率或使用梯度检查点。如果 MACs 大但内存够考虑用更快的算法或硬件加速。这套思路在边缘部署、云端训练、模型选型时都用得上。算清楚这三笔账很多看似奇怪的内存问题都会变得有迹可循。