1. 先搞清楚一个反直觉现象模型文件小不代表运行省内存很多人第一次把训练好的卷积网络拿去部署时都会愣一下磁盘上那个.pt或者.onnx文件才十几兆甚至几兆结果一跑起来进程内存直接飙到几百兆甚至上 G。尤其是做嵌入式或者边缘设备的朋友看到这个数字血压直接上来——明明文件这么小凭什么运行时吃这么多内存我自己最早踩这个坑是在一个图像分类的小项目上。模型权重文件 12MBFP32 存储参数量大概三百万出头。我天真地以为推理时内存占用也就十几兆结果用工具一测常驻内存接近 400MB。当时第一反应是是不是哪里内存泄漏了排查了半天才发现根本不是泄漏而是我对模型文件大小和运行时内存占用这两件事的理解从一开始就错位了。这篇文章就是想把这件事彻底讲透。核心就一句话模型文件只存了权重而运行时内存要同时装下权重、激活值、中间缓冲区、框架开销这四类东西。其中真正的大头往往不是权重而是激活值activation和各类临时缓冲区。我会把卷积这块的账拆成三笔来算——参数量、MACs乘加运算次数、激活内存把每一笔的来龙去脉、计算公式、实际影响都摆出来。适合所有做 CNN 部署、模型压缩、边缘推理的工程师也适合刚入门想搞懂内存到底花在哪的深度学习学习者。看完你至少能做到拿到一个模型不用跑就能大致估出它的运行时内存量级知道该从哪个方向去省。2. 把卷积的三笔账拆开参数量、MACs、激活内存2.1 为什么文件大小约等于参数量乘精度先说最基础的一笔账也是大家最容易理解的模型文件里存的主要就是权重和偏置。一个卷积层的参数量公式是参数量 卷积核高 × 卷积核宽 × 输入通道数 × 输出通道数 输出通道数偏置举个具体例子。一个3×3卷积输入 64 通道输出 128 通道那么参数量就是3 × 3 × 64 × 128 128 73728 128 73856个参数。如果按 FP32 存储每个参数 4 字节这一层就是73856 × 4 ≈ 295KB。整个模型如果有几十层这样的卷积加起来就是十几兆到几十兆这跟我们在磁盘上看到的文件大小是对得上的。所以第一笔账的结论很直接文件大小 ≈ 参数量 × 每参数字节数。FP32 是 4 字节FP16 是 2 字节INT8 是 1 字节。这也是为什么量化能把模型文件压到原来的四分之一甚至更小——它改的就是这个每参数字节数。但这里有个关键点必须点破参数量只决定了静态存储它跟运行时内存的关系只是运行时内存里的一小块。很多人算内存就只算这一块然后得出模型才 12MB内存怎么会用 400MB的困惑问题就出在这。2.2 第二笔账MACs 决定的是计算量不是内存但它间接影响内存MACs 是 Multiply-Accumulate 的缩写指乘加运算次数衡量的是计算量。卷积层的 MACs 公式是MACs 卷积核高 × 卷积核宽 × 输入通道数 × 输出通道数 × 输出特征图高 × 输出特征图宽注意它比参数量公式多了输出特征图的高和宽这一项。这意味着什么意味着MACs 跟输入分辨率强相关而参数量跟分辨率无关。同一个卷积层输入 224×224 和输入 448×448参数量完全一样但 MACs 会变成原来的 4 倍。那 MACs 跟内存有什么关系关系是间接但重要的MACs 越大通常意味着特征图越大、中间结果越多而框架为了高效执行会为这些中间结果分配缓冲区。更直接的是很多推理引擎会做算子融合、内存复用但复用的前提是它能算清楚生命周期特征图越大可复用的空间就越难压缩。所以 MACs 本身不是内存但它是内存需求的一个强信号。我一般会用一个粗略的经验关系来快速判断如果 MACs 在 G 级别几十亿次那运行时内存基本跑不掉几百兆。这个经验在 ResNet、MobileNet 这类常见结构上屡试不爽。2.3 第三笔账激活内存才是运行时的大头这是最容易被忽略、但往往最吃内存的一笔。激活值activation指的是每一层输出的特征图它们在前向传播时需要被保留因为反向传播要用训练时或者因为后续层要读推理时。激活内存的估算公式单个特征图单层激活内存 批大小 × 输出通道数 × 输出特征图高 × 输出特征图宽 × 每元素字节数还是用例子说话。假设 batch size 为 1某层输出 256 通道特征图 56×56FP321 × 256 × 56 × 56 × 4 3,211,264 字节 ≈ 3.06MB单层就 3MB。一个网络几十层如果每层激活都保留轻松上百兆。而且训练时因为要存反向传播的中间值激活内存通常是推理时的好几倍。这就是为什么训练比推理吃内存——不是权重变大了是激活值要留的多了。推理时虽然不用存反向的激活但框架仍然需要为当前层的输入输出分配缓冲区加上算子融合不完美、内存池碎片等原因实际占用往往比理论最小值高不少。这笔账才是模型文件小但运行吃内存的真正答案。3. 逐层算账一个具体卷积网络的内存拆解实录3.1 搭一个可复现的小网络来算光讲公式太虚我直接搭一个结构清晰的小 CNN把每一笔账都算给你看。结构如下输入 3×224×224层类型输入通道输出通道核大小步长输出尺寸1Conv3323×3232×112×1122Conv32643×3264×56×563Conv641283×32128×28×284Conv1282563×32256×14×145Conv2565123×32512×7×76FC512×7×71000--1000这个结构参考了经典分类网络的思路通道翻倍、空间减半很适合用来演示内存分布。3.2 参数量逐层计算按公式kh × kw × Cin × Cout Cout逐层算第 1 层3×3×3×32 32 864 32 896第 2 层3×3×32×64 64 18432 64 18496第 3 层3×3×64×128 128 73728 128 73856第 4 层3×3×128×256 256 294912 256 295168第 5 层3×3×256×512 512 1179648 512 1180160第 6 层全连接512×7×7×1000 1000 25088000 1000 25089000总参数量约896 18496 73856 295168 1180160 25089000 ≈ 26,657,576约 2666 万参数。FP32 存储26657576 × 4 ≈ 106.6MB。看到没光全连接层就占了 2500 万参数接近总量的 94%。这也是为什么后来的网络都改用全局平均池化替代大 FC 层——FC 层参数量太恐怖了。这个模型文件大概 100MB 出头符合预期。3.3 激活内存逐层计算现在算激活。推理时 batch1FP32逐层算单层激活内存C × H × W × 4第 1 层输出32 × 112 × 112 × 4 1,605,632 ≈ 1.53MB第 2 层输出64 × 56 × 56 × 4 802,816 ≈ 0.77MB第 3 层输出128 × 28 × 28 × 4 401,408 ≈ 0.38MB第 4 层输出256 × 14 × 14 × 4 200,704 ≈ 0.19MB第 5 层输出512 × 7 × 7 × 4 100,352 ≈ 0.10MB单看一层不大但注意第 1 层输出就有 1.53MB而且实际运行时框架不会只保留一层——它需要同时持有输入、输出、以及可能的临时缓冲。如果框架不做内存复用把所有层激活都留着加起来约1.530.770.380.190.10 ≈ 2.97MB。这个数字看起来不大但这是 batch1 的情况。3.4 把 batch 和训练模式加进来数字立刻变脸关键来了。把 batch size 提到 32训练常见值激活内存直接乘 322.97MB × 32 ≈ 95MB再加上训练时要保存反向传播所需的中间值通常还要再乘 1.5 到 2 倍激活内存轻松到 150MB 以上。这时候权重才 106MB激活已经快追平了。如果输入分辨率翻倍到 448×448激活内存再乘 4直接爆掉。这就是为什么很多人训练时遇到 OOM显存/内存不足第一反应是模型不大啊但真正吃内存的是 batch 和分辨率撑起来的激活值。参数量决定模型文件batch 和分辨率决定运行时内存这句话我建议你记牢。4. 除了三笔账还有哪些隐形内存在偷吃4.1 框架运行时开销别小看这一块前面算的都是理论数据但实际跑起来进程内存总比理论值高。高出来的部分很大一块是框架自身的运行时开销。以常见的深度学习框架为例加载一个模型时框架会加载算子库cuDNN、MKL 等这些库本身就有几十到上百 MB建立内存池memory pool预分配一块内存备用避免频繁申请释放维护计算图、张量元数据、设备上下文等结构我实测过一个空进程加载框架但不跑任何模型内存占用就可能到 100-200MB。这部分是固定成本跟模型大小无关但它是真实存在的。所以当你看到模型 12MB进程 400MB时先扣掉这 100-200MB 的框架底噪剩下的才是模型真正带来的。4.2 内存池与碎片为什么实际占用比理论高现代推理引擎几乎都会用内存池。原理是与其每次算子执行都向系统申请内存不如一次性申请一大块自己管理分配。好处是快坏处是峰值占用会偏高因为池子里的内存即使暂时不用也不会立刻还给系统。更麻烦的是内存碎片。如果特征图大小不一内存池里会出现很多用不上又回收不了的空洞。我遇到过最夸张的情况理论激活内存 80MB实际进程占用 300MB多出来的全是碎片和池子预留。解决办法通常是调整内存池配置或者用支持动态内存复用的推理引擎比如某些引擎的内存复用开关让生命周期不重叠的张量共享同一块内存。4.3 数据类型FP32 换成 FP16/INT8 能省多少这是最直接的省内存手段。把 FP32 换成 FP16权重和激活的内存都直接减半换成 INT8减到四分之一。但要注意几点权重可以离线量化激活通常需要在线量化精度损失要评估不是所有算子都支持低精度混合精度下会有转换开销INT8 量化对激活值的动态范围敏感校准集选不好精度掉得厉害我一般建议先量化权重观察精度如果精度可接受再考虑激活量化。权重 FP16 基本无损是性价比最高的一步。5. 常见问题与排查技巧实录5.1 内存占用排查速查表现象可能原因排查方向进程内存远大于模型文件激活值 框架开销用内存分析工具看张量分配训练时 OOM推理正常batch 太大或激活未释放减小 batch检查是否保留计算图内存随运行时间持续增长内存泄漏检查是否有张量被意外持有峰值内存远高于稳态内存池预留 碎片调整池配置启用内存复用换小模型内存没明显降瓶颈在激活或框架算激活内存看是否分辨率过高5.2 几个我踩过的坑第一个坑以为torch.no_grad()就能省内存。它确实能省反向传播的激活但前向的中间张量该分配还是分配。真正省内存要靠推理引擎的图优化和内存复用。第二个坑忽略输入分辨率的影响。有次我把输入从 224 提到 320模型没变内存直接涨了 2 倍多。因为激活内存跟特征图面积成正比面积是边长的平方。第三个坑在循环里反复创建模型。有次写测试代码每个 batch 都重新 load 一次模型内存池不断新建很快就爆了。正确做法是模型只加载一次复用。5.3 省内存的实操建议推理时用torch.inference_mode()或对应框架的推理模式比no_grad更彻底能用量化就用FP16 优先INT8 谨慎减小 batch或者用梯度累积模拟大 batch输入分辨率按需下调别盲目追求高分辨率选支持内存复用的推理引擎部署时关掉不必要的调试功能6. 把账算清楚之后部署思路就清晰了把这三笔账算明白之后我做模型部署的思路完全变了。以前拿到模型先看文件大小现在先看三样东西参数量决定权重内存、MACs 和输入分辨率决定激活内存量级、以及目标硬件的内存上限。三者一对比能不能跑、要不要量化、batch 设多少心里基本有数。举个最近的例子。有个项目要在内存受限的设备上跑一个分割网络模型文件 40MB看着不大。但我一算激活输入 512×512中间特征图最大到 128×128×256FP32 下单层激活就 16MB加上多尺度输出和内存池峰值轻松过 200MB。设备只有 256MB 可用内存直接告急。后来做了三件事权重 INT8 量化、输入降到 384×384、启用推理引擎的内存复用。最终峰值压到 90MB 左右跑得很稳。所以模型文件小为什么还吃内存这个问题答案从来不是单一的。它是权重、激活、框架开销、内存池、数据类型这几笔账叠加的结果。你把每一笔都算清楚就不会再被文件大小迷惑也能精准地找到省内存的着力点。真正吃内存的往往不是你以为的那个模型而是它跑起来之后撑开的那些中间状态。
阅读完成 · 觉得有帮助?