ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

卷积神经网络内存膨胀:参数量、MACs与运行时内存的深度解析

卷积神经网络内存膨胀:参数量、MACs与运行时内存的深度解析 1. 一个让很多人困惑的现象刚接触深度学习部署的朋友经常会遇到一个让人挠头的问题明明模型文件才几十兆为什么一跑起来内存就飙到几个G我最早做模型推理服务的时候就踩过这个坑一个ResNet-50的权重文件也就98MB左右结果服务一启动内存直接吃掉1.5G当时第一反应是是不是哪里内存泄漏了排查了半天才发现问题根本不在代码而在于我对卷积层的内存开销理解太浅。这个现象其实非常普遍。你拿一个训练好的模型文件看它的大小觉得就这么点东西能占多少内存但实际运行时内存占用往往是文件大小的十几倍甚至几十倍。这中间的差距从哪来的答案就藏在卷积运算的三笔账里参数量、MACs乘加运算次数、以及运行时内存。这三者之间的关系很多人是模糊的甚至混为一谈。这篇文章适合所有做模型部署、推理优化、边缘端落地的同学。不管你是刚入门的新手还是已经做过几个项目但没深究过内存问题的工程师我都会把这三笔账从头算清楚。你需要的基础知识只有一点知道卷积大概是怎么回事知道什么是FP32。其他的我用人话给你讲明白。2. 先把三笔账的概念理清楚2.1 参数量到底算的是什么参数量说白了就是模型里所有需要学习的权重和偏置的总数。对于一个标准的二维卷积层参数量的计算公式是参数量 卷积核宽 × 卷积核高 × 输入通道数 × 输出通道数 输出通道数偏置举个例子一个3×3的卷积层输入通道64输出通道128那么参数量就是 3×3×64×128 128 73,856。这个数字乘以4FP32每个参数占4字节大概是295KB。你看一个卷积层的参数量其实不大。整个模型的参数量加起来就是你在文件里看到的那个大小。比如ResNet-50大约2500万参数乘以4字节差不多100MB跟文件大小对得上。所以参数量决定的是模型文件的大小它跟运行时内存的关系是间接的。2.2 MACs为什么才是计算量大头MACs全称是Multiply-Accumulate operations也就是乘加运算次数。每一次卷积操作本质上就是做一次乘法和一次加法。MACs衡量的是模型的计算量也就是这个模型跑一次要算多少次。还是那个3×3卷积输入64通道输出128通道假设输入特征图是56×56那么MACs 3×3×64×128×56×56 ≈ 2.3亿次这个数字就比参数量大得多了。参数量是7万多MACs是2.3亿差了三千倍。这就是为什么模型文件小但计算量大的原因——参数量决定存储MACs决定计算。很多人会把MACs和FLOPs搞混。FLOPs是浮点运算次数一次MAC算两次FLOP一次乘、一次加所以FLOPs通常是MACs的两倍。但在实际工程中大家更习惯用MACs因为它更直接地反映了硬件要做多少次乘加。2.3 运行时内存到底吃在哪里运行时内存这才是真正让内存飙升的元凶。它主要包括三部分权重内存所有参数加载到内存里这部分跟参数量直接相关但通常不是大头。激活内存每一层的输出特征图都要存在内存里因为反向传播或者某些推理框架需要保留中间结果。这部分跟特征图大小直接相关。工作内存卷积运算过程中硬件或框架需要额外的缓冲区来做im2col、矩阵乘法、临时存储等操作。激活内存往往是最容易被忽略的。一个56×56×128的特征图FP32下就是56×56×128×4 1.6MB。看起来不大但一个ResNet有几十层每层都存一份加起来就是几十MB。如果是高分辨率输入比如512×512那特征图直接翻几十倍激活内存轻松上G。注意推理时如果框架做了内存复用比如TensorRT、ONNX Runtime的memory pool激活内存可以大幅降低。但训练时因为要反向传播几乎所有中间激活都得留着内存占用会更高。3. 参数量、MACs、内存三者的真实关系3.1 为什么参数量小不代表内存小这是最核心的认知误区。参数量小只说明模型文件小但运行时内存取决于同时活跃的张量数量。一个模型可能参数量只有几百万但如果它的特征图很大激活内存就会很高。举个极端的例子一个1×1卷积输入通道1输出通道1参数量只有2个。但如果输入特征图是4096×4096那么输出也是4096×4096激活内存就是4096×4096×4 64MB。参数量2个内存64MB差距三万倍。所以你在部署模型时不能只看模型文件大小必须看输入分辨率和特征图通道数。这两个才是决定运行时内存的关键。3.2 MACs和内存的关系MACs高通常意味着计算量大但不一定内存就高。比如深度可分离卷积MACs比标准卷积低很多但内存占用可能差不多因为特征图大小没变。反过来有些操作MACs不高但内存占用很高。比如concat操作它本身不做乘加但需要把两个特征图拼在一起内存直接翻倍。再比如上采样MACs几乎为零但输出特征图是输入的4倍2倍上采样内存直接涨4倍。所以MACs和内存是两个独立的维度优化的时候要分开看。你不能说我MACs降下来了内存就一定降下来了这是两码事。3.3 一个具体的计算示例我拿一个实际的卷积层来算一遍让你有个直观感受。假设输入特征图224×224×64卷积核3×3输出通道128stride1padding1。参数量3×3×64×128 128 73,856FP32下约288KB。MACs3×3×64×128×224×224 ≈ 37亿次。激活内存输出特征图224×224×128FP32下224×224×128×4 25.7MB。如果框架需要保留输入和输出那就是25.7 12.8 38.5MB。你看参数量288KB激活内存38.5MB差了130多倍。这就是为什么模型文件小但内存吃得多。如果输入分辨率变成448×448激活内存直接变成原来的4倍154MB。如果通道数再翻倍那就是308MB。一层就这么多几十层加起来内存不上G才怪。4. 卷积内存膨胀的底层原理4.1 im2col带来的内存放大很多推理框架在实现卷积时会用im2colimage to column的方法。简单说就是把卷积操作转换成矩阵乘法。具体做法是把输入特征图中每个卷积窗口覆盖的区域拉成一列形成一个矩阵然后用这个矩阵跟卷积核矩阵做乘法。这个方法的优点是可以用高度优化的矩阵乘法库比如BLAS计算效率高。但缺点是内存放大。一个3×3卷积im2col之后输入矩阵的行数变成原来的9倍因为每个位置要拉出9个元素。虽然列数减少了但总体内存占用还是增加了。具体来说输入特征图224×224×64im2col之后变成(224×224)×(3×3×64) 50176×576的矩阵FP32下就是50176×576×4 115MB。而原始输入只有224×224×64×4 12.8MB。放大了9倍。这就是为什么很多框架在推理时内存飙升——im2col的临时缓冲区太大了。后来大家用Winograd、FFT等算法来减少计算量和内存但im2col仍然是最通用的方法。4.2 特征图的生命周期管理推理框架通常会做内存复用也就是一块内存用完了后面的层可以接着用。但这个复用是有条件的只有当两个张量的生命周期不重叠时才能复用同一块内存。在推理时因为不需要反向传播每一层的输入用完就可以释放所以内存复用效率很高。但在训练时每一层的输入都要留着给反向传播用内存复用效率就很低。这也是为什么同一个模型训练时内存占用可能是推理时的3-5倍。很多人拿训练时的内存需求去估算推理时的内存结果发现推理时内存小很多就是这个原因。4.3 框架层面的内存池机制主流推理框架TensorRT、ONNX Runtime、TVM等都有自己的内存池机制。它们会在初始化时一次性申请一大块内存然后自己管理分配和释放避免频繁调用系统malloc/free。这个机制的好处是减少内存碎片提高分配效率。但缺点是内存占用是峰值决定的。也就是说即使某一时刻内存用量很低内存池也不会把内存还给系统而是留着给后面的层用。所以你看到的内存占用往往是整个推理过程中内存用量的峰值。我实测过一个MobileNetV2模型文件14MB推理时内存峰值大约300MB。其中权重只占14MB剩下的全是激活内存和im2col缓冲区。如果用TensorRT优化内存可以降到150MB左右因为它做了层融合和内存复用。5. 实操如何准确估算和优化内存5.1 手算内存占用的方法如果你想在部署前估算内存占用可以按这个步骤来列出所有层把模型的每一层列出来包括卷积、BN、ReLU、池化等。计算每层输出特征图大小根据输入分辨率、卷积核大小、stride、padding算出每层输出的宽高和通道数。计算激活内存每层输出特征图大小 × 4字节FP32。如果是推理通常只需要保留当前层和下一层的输入所以可以只算峰值。计算im2col缓冲区对于每个卷积层输入特征图大小 × 卷积核面积 × 4字节。加上权重内存所有参数 × 4字节。取峰值把上面所有加起来取整个推理过程中的最大值。这个估算方法不是100%准确因为框架的具体实现会有差异但能给你一个数量级的判断。我一般会在这个基础上乘以1.5作为安全余量。5.2 用工具实测内存占用手算毕竟麻烦实际项目中我更多用工具来测。常用的方法有PyTorch用torch.cuda.max_memory_allocated()看GPU内存用tracemalloc看CPU内存。ONNX Runtime开启profiling会输出每个节点的内存占用。TensorRT用trtexec工具加--dumpProfile参数可以看到每层的内存和时间。系统工具Linux下用/usr/bin/time -v看峰值内存或者用ps、top实时监控。我一般会先用工具测出峰值内存然后跟手算结果对比看看差距在哪里。如果差距很大通常是某个层的im2col缓冲区特别大或者框架没有做好内存复用。5.3 降低内存的几种实用手段如果你发现内存占用太高可以尝试这几种方法降低输入分辨率这是最直接有效的。输入从224降到160激活内存直接降到原来的51%。精度可能会掉一点但很多时候可以接受。使用FP16或INT8FP32换成FP16内存直接减半。INT8再减半。现在很多硬件都支持FP16和INT8加速精度损失可以通过量化感知训练来弥补。层融合把ConvBNReLU融合成一个操作减少中间特征图的存储。TensorRT和ONNX Runtime都支持这种优化。内存复用确保框架开启了内存复用。PyTorch的torch.no_grad()、ONNX Runtime的enable_mem_pattern、TensorRT的--memPoolSize都可以控制。换算法用Winograd代替im2col可以减少计算量和内存。但Winograd对硬件有要求不是所有平台都支持。提示优化内存时不要只盯着一个手段通常是组合使用。比如先降分辨率再上FP16再做层融合三管齐下内存可以降到原来的1/4甚至更低。6. 常见问题与排查技巧6.1 为什么模型文件小但内存大这个问题前面已经解释过了核心原因是激活内存和im2col缓冲区。模型文件只包含参数量而运行时内存还包括特征图、临时缓冲区等。一个经验法则是运行时内存通常是模型文件的10-30倍。如果你的模型文件是100MB那运行时内存1-3G是正常的。6.2 内存泄漏还是正常占用很多人看到内存高就怀疑是内存泄漏。区分方法很简单看内存是否持续增长。如果内存稳定在一个峰值不再涨那就是正常占用。如果内存一直涨跑几次就OOM那才是泄漏。内存泄漏常见的原因有没有用torch.no_grad()、没有释放中间变量、DataLoader的worker没关等。我踩过的一个坑是在循环里不断创建新的tensor但没有释放导致内存一直涨。后来改成复用tensor问题就解决了。6.3 不同框架的内存表现差异同一个模型在不同框架下内存占用可能差很多。我实测过一个模型框架内存峰值说明PyTorch (eager)1.2GB默认不做优化内存最高PyTorch (JIT)800MB做了图优化内存降低ONNX Runtime600MB内存复用做得好TensorRT400MB层融合FP16内存最低这个差异主要来自框架的优化程度。TensorRT因为做了层融合和FP16内存最低。ONNX Runtime次之。PyTorch eager模式最费内存但最灵活。6.4 常见问题速查表问题可能原因解决方法内存占用是模型文件的20倍激活内存im2col缓冲区降低分辨率、用FP16、层融合内存持续增长不释放内存泄漏检查no_grad、释放中间变量推理时内存比训练时小很多训练要保留激活做反向传播正常现象推理时内存复用效率高换了框架内存变化很大框架优化程度不同选优化好的框架如TensorRT某层内存特别大im2col缓冲区大换Winograd或FFT算法6.5 几个容易踩的坑坑一只看模型文件大小估算内存。这是最常见的错误。模型文件只是参数量运行时内存还包括激活和缓冲区。我一般会按模型文件的15-20倍来估算。坑二忽略输入分辨率的影响。输入分辨率翻倍激活内存翻4倍。很多人调模型时只关注精度忽略了分辨率对内存的影响。坑三以为FP16只是加速。FP16不仅加速还能直接减半内存。如果你的硬件支持FP16强烈建议用。坑四不做内存复用。有些框架默认不开内存复用需要手动开启。比如ONNX Runtime的enable_mem_patternTensorRT的--memPoolSize。坑五忽略batch size的影响。batch size翻倍激活内存也翻倍。推理时如果不需要batch就设成1。7. 一个完整的计算实例我拿一个实际的模型来算一遍让你有个完整的感受。假设我们有一个简单的CNN输入224×224×3结构如下Conv1: 3×3, 3→64, stride1, padding1Conv2: 3×3, 64→128, stride1, padding1Conv3: 3×3, 128→256, stride1, padding1FC: 256×7×7→1000参数量Conv1: 3×3×3×6464 1,792Conv2: 3×3×64×128128 73,856Conv3: 3×3×128×256256 295,168FC: 256×7×7×10001000 12,545,000总计约12.9M参数FP32下约51.6MBMACsConv1: 3×3×3×64×224×224 ≈ 8.7亿Conv2: 3×3×64×128×224×224 ≈ 37亿Conv3: 3×3×128×256×224×224 ≈ 148亿FC: 256×7×7×1000 ≈ 12.5亿总计约206亿MACs激活内存假设只保留当前层输出Conv1输出224×224×64×4 12.8MBConv2输出224×224×128×4 25.7MBConv3输出224×224×256×4 51.4MB峰值51.4MBim2col缓冲区以Conv3为例输入224×224×128im2col后224×224×3×3×128×4 231MB你看im2col缓冲区231MB比激活内存51.4MB大得多。这就是为什么实际运行时内存会飙升到几百MB甚至上G。如果换成FP16所有数字减半。如果输入分辨率降到112×112所有数字降到1/4。如果做层融合im2col缓冲区可以省掉。组合使用内存可以从几百MB降到几十MB。8. 我个人的一些经验体会做了这么多年的模型部署我最大的体会是不要用模型文件大小来估算内存。这个习惯害了很多人。我现在拿到一个新模型第一件事是看输入分辨率和特征图通道数这两个决定了内存的下限。然后看框架的优化程度这决定了内存的上限。另一个体会是内存优化没有银弹。降分辨率、FP16、层融合、内存复用每个手段只能解决一部分问题必须组合使用。我一般会先做层融合和内存复用这两个不需要改模型效果也最明显。然后再考虑降分辨率或量化这两个会影响精度需要权衡。最后说一个容易被忽略的点batch size对内存的影响是线性的。很多人推理时习惯用batch size8或16觉得能提高吞吐。但如果内存不够batch size1反而更稳。我一般会先测batch size1的内存然后根据余量决定能开到多大。还有一个坑是不同硬件的内存表现不一样。同样的模型在服务器CPU上跑和在边缘设备上跑内存占用可能差很多。边缘设备通常内存更紧张优化要更激进。我一般会在目标硬件上实测而不是在开发机上估算。如果你也在做模型部署建议你养成一个习惯每次拿到新模型先算一遍参数量、MACs和激活内存心里有个数。然后用工具实测对比一下。时间长了你就能凭经验判断一个模型大概吃多少内存这对选型和优化非常有帮助。
返回列表