卷积算子优化全攻略：从im2col到稀疏化

最近在折腾一个推理引擎的卷积算子从最朴素的循环实现一路优化到寄存器级中途还把剪枝稀疏化的路子走了一遍。回过头看卷积运算的高效实现这件事绝不是背几个优化口诀就能搞定的它横跨算法、数据结构、硬件架构好几个层面每一步取舍都是权衡。这篇内容就是把我自己实践下来的经验做一个系统梳理从卷积的计算本质开始到im2col、Winograd这些算法级优化再到GPU/CPU上的访存与向量化落地最后聊目前很受关注的稀疏化路线。希望能给正在做深度学习推理加速、算子优化或者底层系统开发的同行一些可直接参考的素材。如果你刚接触这块建议先把第1节的基础吃透后面每一节都是在前面的复杂度分析上逐步叠加思路如果你已经做过一些算子开发可以直接跳到第4、5节看稀疏化的坑那部分踩坑实录我觉得比任何论文里的加速比表格都有参考价值。1. 先把卷积的计算本质掰清楚1.1 卷积的两种理解方式滑动窗口与矩阵乘法卷积的数学定义我就不抄教科书了直接说工程上最常用的两种理解。第一种是滑动窗口视角也是初学者最容易接受的方式卷积核在输入特征图上按步长滑动每个位置做一次逐元素相乘再求和。这种视角的好处是直观坏处是容易让人误以为卷积就是“嵌套循环”从而低估它的优化空间。实际上循环只是最原始的描述方式不是高效实现的方式。第二种是矩阵乘法视角把卷积操作映射成矩阵乘法。这个视角是im2col等一系列高效算法的出发点。什么意思呢就是把输入特征图上每个滑动窗口对应的数据块取出来拉平成一行所有窗口的行堆叠成一个大矩阵卷积核展开成另一个矩阵两者做矩阵乘法结果就是输出特征图。矩阵乘法是底层被优化到极致的运算BLAS库、GPU的Tensor Core都针对它做了大量优化所以把卷积转化成矩阵乘法本质上是“蹭”矩阵乘法的优化红利。两种视角没有谁对谁错我个人的习惯是分析计算量、推导复杂度时用滑动窗口视角设计实现和优化时用矩阵乘法视角。搞清楚自己当前处于哪个阶段就知道该用哪种视角去思考。1.2 一次普通卷积到底有多少计算量先说一个很多初学者容易忽略的点FLOPs和FLOPS是两个完全不同的概念。前者是浮点运算次数Floating Point Operations后者是每秒浮点运算能力Floating Point Operations Per Second。神经网络里常说的FLOPs其实就是一次前向推理需要做多少次乘加运算。来算一笔具体的账。假设输入是224×224分辨率、64个通道的特征图卷积核是3×3、输出通道也是64。那么一次卷积的计算量可以按如下方式估算输出特征图的空间尺寸是224×224假设padding1且stride1保持尺寸不变每个输出位置要做3×3×64 576次乘加运算。所以单个输出通道的FLOPs是224×224×576约2890万。总共有64个输出通道总FLOPs就是224×224×576×64约18.5亿次乘加。一次普通的3×3卷积就是18.5亿次运算这还只是单层。实际模型动辄几十上百层累计起来就是几百亿甚至上千亿次运算。这就是为什么卷积优化如此重要的根本原因任何微小的单次运算效率提升放大到整个网络规模之后都会变成肉眼可见的性能差异。我在实际调优的时候经常跟团队说“不要小看哪怕是5%的算力节省放到推理侧可能就是每秒几十帧的差距”。2. 算法层面的优化怎么让卷积算得更快2.1 im2col把卷积“铺开”成矩阵乘法im2col是“image to column”的缩写核心操作就是把输入特征图按滑动窗口重新排列成矩阵。这个方案在深度学习框架里非常普遍早期的Caffe基本就是靠它把卷积接入BLAS库的。具体流程可以这样理解假设原始输入是C×H×W的三维张量im2col之后会变成(K×C×kh×kw)×(H_out×W_out)的二维矩阵其中K是卷积核个数kh和kw是卷积核尺寸H_out和W_out是输出特征图的尺寸。每一列对应一个输出位置的感受野数据。卷积核则展开成K×(C×kh×kw)的矩阵。两个矩阵相乘得到K×(H_out×W_out)的输出再reshape成K×H_out×W_out就是输出特征图。直接说这个方案的优缺点。优点是实现简单、矩阵乘法高度优化性能稳定可预测缺点是显存或内存占用会显著增加。还是用上面的例子输入是64×224×224im2col之后矩阵的列数是输出位置的个数224×22450176行数是64×3×3576那么新矩阵的规模就是576×50176约2880万个元素。相比原始输入的320万个元素内存扩张了将近9倍。这还只是一层所以im2col在硬件资源紧张的端侧场景要慎用但在GPU和服务器端是可以接受的因为GPU显存带宽高、容量大BLAS库的矩阵乘法效率能抵消内存扩张带来的开销。我在实践中的经验是im2col适用于卷积核尺寸较大比如5×5、7×7或通道数很多的场景因为这些情况下矩阵乘法的计算密度高内存扩张的代价相对可控。对于1×1卷积im2col退化成纯粹的矩阵乘法不需要做任何重排直接拿输入当矩阵乘就行。2.2 Winograd用加法换乘法的经典套路如果说im2col是“空间换时间”那Winograd就是“加法换乘法”。它的核心思想是在卷积的计算过程中存在大量可以复用的中间结果通过数学变换减少乘法次数因为乘法比加法更耗时。这里用最简单的F(2,3)举例也就是输出2个点、卷积核长度为3的一维卷积。原始计算需要2×36次乘法Winograd变换之后输入变换需要3次加法输出变换需要4次加法中间只有4次乘法。乘法次数从6降到了4减少了1/3。二维卷积可以把一维变换组合起来。对3×3卷积核和2×2输出块也就是F(2×2, 3×3)乘法次数会从36次降到16次理论加速比是2.25倍。这也是为什么3×3卷积在Winograd下收益特别明显的原因——3×3是深度学习模型中出现频率最高的卷积核尺寸。但Winograd有几个工程上必须注意的坑。第一数值精度会下降。因为变换矩阵里出现了1/2、1/4这样的系数误差会被放大对FP16或INT8精度的模型尤其敏感严重时可能让精度掉得不可接受。第二它适合小卷积核、小输出块但不适合大卷积核。卷积核一大变换矩阵本身就变得复杂中间计算量反而上升。第三实现复杂度高尤其要和硬件访存模式配合时没有经验的人写出来可能比朴素卷积还慢。我见过不止一个团队在Winograd上投入几周最终性能没有超过优化过的im2col就是因为Winograd的访存局部性特别难把控。我的建议是如果你用的是英伟达GPU成熟的推理框架基本已经内置了Winograd实现直接配置开启即可如果要自己在端侧芯片上实现建议先确认芯片的乘法器和加法器时延比例如果这个比例不大Winograd的理论收益本身就会打折。2.3 FFT一条“绕远路”的频域方案还有一种思路是FFT把卷积从空间域转换到频域。根据卷积定理空间域的卷积等于频域的乘法所以可以用FFT把输入和卷积核都变换到频域做逐元素乘法再逆变换回空间域。FFT的优势在于复杂度低。空间域直接卷积复杂度是O(N²)其中N是特征图尺寸而FFT的复杂度是O(NlogN)。当卷积核特别大、输入特别大时比如医学影像处理中的大核卷积FFT的优势就体现出来了。这也是传统图像处理里高斯模糊等操作常用FFT加速的原因。但在深度学习场景中FFT并不常用原因是主流模型使用的卷积核几乎都是3×3最大也不过7×7核尺寸很小FFT的logN开销优势根本体现不出来反而要额外承担两次正变换和一次逆变换的开销。我在处理一个图像处理Demo时试过在小卷积核上用FFT结果性能比朴素卷积还差就果断放弃了。FFT在深度学习里的归宿我觉得更多是图像预处理和一些特殊的频域模型结构。做算子优化的人知道有这么个方案就行别在常规卷积上浪费时间去调它。3. 硬件层面的落地从循环代码到处理器指令3.1 GPU上卷积的核心矛盾访存还是计算GPU的浮点计算能力远超访存带宽所以对GPU来说多数算子的瓶颈不是计算而是数据搬运。卷积的重心在于如何让数据在寄存器、共享内存、全局内存之间流动得更高效。这个领域有个经典问题叫“内存墙”计算单元空转等待数据的时钟周期远大于真正运算的时钟周期。最直接的应对策略是提高数据的复用率。卷积天然具备两种复用输入特征图的复用同一个输入像素会被多个输出位置用到和卷积核的复用同一个卷积核权重会被整个特征图用到。高水平的GPU卷积实现核心就是在调度计算时最大化这两种复用。实际工程中最常见的方法是分块tiling。把输出特征图分成若干个小块让每个线程块处理一小块输出同时把这一小块需要用到的输入数据和卷积核权重加载到共享内存反复使用。这就是类似“手艺人把工具都放在手边不用每次干活都起身去拿”的思路。分块尺寸的选择是个经典调参问题分块太小共享内存带宽利用不充分分块太大共享内存装不下还得处理边界。我在一个模拟项目X里做过GPU卷积的开发发现一个反直觉的经验大通道数场景下仅依赖共享内存缓存卷积核权重就能获得相当大的加速因为权重复用率极高反而缓存输入特征图的分块策略对性能提升没那么显著因为输入数据每次滑动窗口之间重合度不高缓存命中率有限。3.2 CPU上的卷积优化向量化与缓存友好CPU场景更复杂的点是既有AVX这类SIMD向量指令又有层级化的缓存结构任何一环没用好性能都会大打折扣。首先是SIMD向量化。CPU一次可以对多个数据执行同一操作比如AVX512能一次处理16个FP32数据。要让代码吃满SIMD数据布局必须连续且对齐。这也是为什么很多高性能实现在计算时会把通道维放到最内层循环因为同一空间位置上的多个通道在内存中的地址是连续的可以一次性load进向量寄存器。其次是缓存友好。CPU的L1缓存一般只有几十KBL2也就几百KB一旦数据超过缓存容量就会不断触发缓存未命中性能断崖式下跌。卷积的实现要尽量保证数据的访问局部性——即短时间内反复访问的数据尽量在缓存里待着。实际做法就是循环重排序把循环顺序从“遍历所有输出位置→遍历所有通道”改成“分块遍历输出位置→遍历通道”让一小块输入数据在L1/L2缓存中反复命中。我当时在调一个端侧CPU的卷积算子时深刻体验过“一改循环顺序性能翻倍”的现场。原始代码按输出位置顺序遍历每次都需要重新从内存里拉取输入数据改成按通道内小块遍历之后L1缓存命中率从不到40%提升到接近90%单算子耗时直接降了一倍多。缓存优化就是这么立竿见影。3.3 数据布局NCHW、NHWC与更细致的排布数据布局看似是小事其实对性能影响非常大。NCHW通道在前、NHWC通道在后是两种主流格式各有利弊。NCHW的优点是通道维连续方便按通道处理数据在PyTorch默认训练格式中常见缺点是空间上同一位置的不同通道数据被分散在很远的内存地址访问起来不连续。NHWC把通道维放在最后让同一位置的所有通道在内存里连续排布。这种布局对SIMD向量化非常友好CPU可以直接load连续的一段内存做向量运算。这也是TensorFlow和很多移动端推理框架默认使用NHWC的原因。除此之外GPU上还常使用更复杂的排布比如把通道分成多个组组内仍是CHW顺序组间交错布置目的是让不同线程访问的数据尽量连续。实际开发中如果你要调的卷积性能始终上不去建议先检查数据布局是否匹配硬件特性。我踩过最大的坑就是用NCHW格式直接写SIMD代码换来换去发现无论如何都向量化不起来改成NHWC后代码量没增加性能却直线上升。3.4 算子融合从算法到硬件的“最后一公里”最后说一个工程性很强、效果极大的优化手段算子融合。通俗讲就是把多个相邻操作合并成一个操作省掉中间的显存读写。比如卷积之后通常接ReLU激活函数经典做法是卷积写出一份中间结果到显存ReLU再读出来算一次再写回显存。合并之后卷积算出一个值直接流式地做一次ReLU判断再写入最终位置中间那份临时数据直接不存在了。可别小看这一步。每个算子的读写都是全特征图数据量的模型越深中间张量越大省下来的带宽就越可观。我做过统计在某个图像分类模型上仅做ConvReLU融合单层推理耗时能降低10%到15%。如果再往上叠加ConvBatchNormReLU的融合整体减少的中间读写次数非常可观。融合的原理本质上是把“算完存起来、再读出来算”的模式改成“在寄存器或缓存里算完直接往下传”的流水线模式。这也是典型的从算法流程倒推硬件实现的思维方式算法上描述了数据转换的步骤硬件实现上要让数据尽量留在离计算单元近的地方。做算子融合时顺序不能乱比如BatchNorm在推理时可以提前折叠进卷积的权重和偏置里变成纯卷积ReLU的组合这样融合的操作更干净。4. 稀疏化一条看似颠覆、实则妥协的革命4.1 稀疏化为什么有效权重冗余与模型容量稀疏化的思路很直接既然模型参数那么多其中必然有大量的接近零或者不重要的权重把这些权重置零模型变小、计算量降低精度却未必明显下降。这个想法的底气来自深度神经网络的过度参数化特征。一个训练好的模型很多神经元的连接权重对最终输出贡献很小可以安全地移除。稀疏化本质上是在模型容量和计算代价之间做权衡保留容量去掉冗余。为什么说“看似颠覆、实则妥协”因为稀疏化并没有改变卷积的计算本质它只是承认了一个事实并不是所有的乘法运算都同等重要。如果把卷积的所有乘法都一视同仁就必须全部算完如果接受了“很多权重是零”这一事实就有机会跳过大量无效计算。但稀疏化有一个根本矛盾跳过零值的计算方法和硬件本身“均匀计算”的设计哲学并不完全兼容。GPU和CPU本质上都是均匀的向量处理器硬件不会天然知道哪些权重是零、哪些不是。稀疏化需要额外的元数据来描述“哪里稀疏、哪里不稀疏”这部分开销如果设计不好省下的计算量会被额外的访存和控制逻辑吃掉。这也是稀疏化从论文到工程落地之间最大的鸿沟。4.2 剪枝完整流程训练、剪枝、微调三步走实际落地中一个人最常用的稀疏化方法是权重剪枝整个流程可以总结为三个阶段。第一阶段是正常训练得到一个收敛的稠密模型。这个阶段的目标是让模型精度足够高、权重分布足够稳定。很多人在这个阶段就犯错了还没训练好就急着剪枝结果后续精度怎么都补不回来。第二阶段是剪枝。对每个权重按绝对值大小排序把绝对值小于阈值或者排在末尾一定比例的权重直接置零。最简单的是全局阈值法——设定一个全局阈值低于阈值的权重全部置零。常用的还有按层阈值法——每层独立设置阈值保证每层的稀疏率一致。我在实践中更推荐按层阈值法因为它能避免某些层的权重被全部剪光、另一些层却几乎没有稀疏的情况。第三阶段是微调。剪枝后的模型直接在原数据集上继续训练几个epoch让剩余权重适应被移除的路径。这一步很关键除非你的稀疏率极低比如低于20%否则不微调的模型精度几乎必然掉点。微调时建议使用较小的学习率大约是原始训练学习率的十分之一到五分之一。整个流程中最反直觉的经验是一次性剪到目标稀疏率比逐步剪更好。我试过“每次剪5%剪到50%再微调”的渐进方案结果精度和一次性剪到50%差不多但花费的训练时间却翻了好几倍。所以除非你需要监控精度下降曲线来选阈值否则直接一次剪到位再微调是更经济的选择。4.3 结构化稀疏与非结构化稀疏硬件友好度的分水岭剪枝得到的稀疏模式可分为非结构化和结构化两种这两种的工程价值有天壤之别。非结构化稀疏就是权重矩阵中随意散布着零值每个位置独立决定是否置零。这种模式能获得很高的稀疏率比如90%以上但硬件很难加速。因为计算时无法用连续的向量指令批量跳过随机位置的零值你需要逐个检查权重是否为非零才能决定是否做乘法。我在CPU上做过测试非结构化稀疏率80%时提速比例远低于理论值有时候甚至比稠密计算更慢因为稀疏判断和索引开销反而更大。结构化稀疏是约束稀疏位置的分布模式让零值出现在规则的位置上。常见的有行稀疏、列稀疏、块稀疏。比如块稀疏把权重矩阵划分成固定大小的块比如2×4每个块要么整块保留要么整块置零。这种模式可以用一条位图元数据描述整个块的稀疏状态硬件一旦识别块是零可以直接跳过整块计算不需要逐个元素判断。从实际的工程选择角度来看如果是服务器端GPU推理可以接受非结构化稀疏配合专门的稀疏计算库因为GPU并行度高可以在一定程度上掩盖稀疏判断的开销如果是移动端或嵌入式芯片建议只考虑结构化的稀疏方案否则无法利用现有的SIMD和缓存优化。4.4 稀疏格式的底层原理CSR与CSC要真正落地稀疏化绕不开稀疏矩阵存储格式最基础的是CSRCompressed Sparse Row和CSCCompressed Sparse Column。CSR的核心思路是只存储非零元素再加一份行偏移数组和列索引数组。比如一个4×4的矩阵如果只有6个非零元素稠密存储需要16个元素CSR只需要6个数值6个列索引4个行偏移通常行偏移长度是行数1。稀疏率越高节省的存储空间越大。CSC是CSR的转置版存储时优先按列压缩适合按列访问的场景。普通卷积权重在内存中按输出通道或输入通道排列选择CSR还是CSC取决于后续计算时按行还是按列访问数据。稀疏矩阵的乘法逻辑和稠密乘法完全不同。稠密矩阵乘法是固定的三重循环稀疏矩阵乘法需要先解析行偏移找到某一行中所有非零元素再逐个做乘加。这个过程是串行的而且数据访问带有随机性很难向量化。这也是为什么很多稀疏方案最终要配合硬件专门的稀疏单元比如某些AI芯片上设计了稀疏感知的计算单元能在硬件层面直接跳过零块才真正让稀疏化发挥出性能潜力。如果你用的平台没有这类硬件支持建议对稀疏方案保持谨慎先做原型验证再决定是否全面铺开。5. 稀疏化落地路上的坑与解决办法5.1 坑一访存不规则带来的“稀疏惩罚”前面说过非结构化稀疏在硬件上很难提速最核心的问题就是访存不规则。稠密计算时数据是连续排列的预取器可以自动把后续数据提前搬进缓存稀疏数据则是跳着存的无法预取每次访问都可能是缓存未命中。我在一个模拟项目X中实测过一组50%稀疏的矩阵乘法从计算量上看应该比稠密快一倍但实际运行时间只快了两成。剖析之后发现时间几乎都花在索引运算和缓存未命中上ALU反而在空等。遇到这种问题解决方案有两种。一种是改成块稀疏让稀疏模式规则化恢复一定的访存连续性另一种是重新设计数据布局让稀疏矩阵按“非零元素块”存储而不是按行散存。后者实现复杂但保留了非结构化稀疏的高压缩率优势。5.2 坑二负载不均衡GPU和CPU的并行计算都依赖工作负载均匀分配。稠密卷积的每个输出位置计算量基本相同可以把任务切成大小一致的块分给不同线程核。稀疏卷积则不然有些输出位置对应的权重块全是零一瞬间就完成有些位置几乎全是非零需要大量计算。如果这些位置被分配到不同线程就会出现有的线程早早干完闲着有的线程忙得没完没了。我踩过这个坑在某AI芯片上跑50%稀疏的卷积加速比只有理论值的三成。查下来发现是两个线程组负载严重不均衡其中一个线程组几乎承担了80%的计算量。解决思路通常是重排任务分配把稀疏率相近的输出位置聚在一起或者在任务分配时让每个线程组拿到的稀疏率大体均衡。这里的经验是先统计稀疏率的分布再决定如何切块不要直接沿用稠密版本的任务切分方式。5.3 坑三精度与稀疏率的平衡剪枝一定掉精度只是程度不同。需要反复试不同的稀疏率才能找到一个既能压缩又不至于让精度崩掉的平衡点。实操时我有一个习惯先做小规模实验。在几百张图上跑一次完整剪枝微调实验画出“稀疏率-精度”曲线。一般规律是稀疏率低于60%时精度下降在1%以内超过70%开始明显掉点超过85%精度可能直接跌破可用线。当然不同模型差异很大同一个模型不同层也是不同的。剪枝时按层选择稀疏率也很重要。浅层负责提取低阶特征过于稀疏会影响后续所有层深层特征相对抽象可以更激进地剪。反过来按全局阈值剪往往会发现第一层被剪得稀烂后续层几乎没动模型精度直线下降。5.4 坑四工具链不配合就算算法和稀疏模式设计得再好如果推理框架和底层库不支持稀疏计算一切白搭。很多训练框架导出模型时会保留mask但不会真正移除权重导出的模型文件大小没变小推理时也没有任何加速等于剪了个寂寞。正确的导出流程应该是剪枝后先把权重矩阵转成稀疏格式再重新序列化模型才能确保推理引擎按稀疏模式加载。另一个容易踩的坑是量化与稀疏的叠加顺序。如果你既想剪枝又想量化建议先量化再剪枝。因为量化过程会引入噪声如果先将权重剪没再量化剩余权重的数值分布会被压缩得很窄量化误差会更大。我做过一个图像分类模型的实验先剪枝后量化精度掉了3%以上反过来先量化后剪枝精度只掉了不到1.5%。6. 不同场景下的方案选型与调优参考6.1 几张能直接用来做技术决策的参考表我把自己平时做卷积优化时的选型经验整理成了表格方便直接对照参考。场景推荐优化方向注意事项GPU服务器端推理im2col Tensor Core矩阵乘法内存充足适合大batchGPU服务器端推理小batchWinograd注意精度损失建议FP32下使用端侧CPU推理NHWC布局 SIMD向量化 算子融合重点投入缓存友好优化端侧NPU/DSP结构化稀疏 量化需硬件支持稀疏计算否则别上稀疏大卷积核或图像预处理FFT核尺寸小于7时收益为负存储受限场景非结构化稀疏 CSR压缩推理时追求压缩率不追求计算加速这个表格是一个基础框架实际项目一定需要结合硬件特性调整。做技术选型的时候我建议用“一个算子两个API”的思路来验证先实现一个朴素但绝对正确的版本作为baseline再实现候选优化版本两个版本都跑性能和精度对比数据说话而不是拍脑袋猜。6.2 性能剖析的通用方法优化的前提是定位瓶颈。我自己固定用以下三个步骤定位卷积算子的性能瓶颈第一步确认是访存瓶颈还是计算瓶颈。可以通过计算当前算子的实际FLOPs和有效带宽占用对比硬件峰值来快速判断。如果有效带宽已经接近峰值访存是瓶颈如果FLOPs接近峰值计算是瓶颈。这两种情况的优化方向完全不同。第二步看缓存命中率和指令流水线停顿。缓存命中率低说明数据布局或分块策略有问题流水线停顿多说明代码依赖关系太重并行度不够。第三步做“拆解计时”。把卷积拆成数据重排、矩阵乘、输出写入几个子阶段分别计时。这个方法特别有用可以精确定位时间到底花在哪一步——很多时候你会发现矩阵乘法本身没多少耗时反而是im2col的数据重排过程成了主要开销。6.3 我的几条切身体会做卷积优化这几年我的体会有几条特别想说。一是“正确的朴素实现”比“没有实现”更值钱。团队里有新人一上来就想用顶尖优化结果写了两周代码连正确性都没验证这种方向错了。先实现最笨、最正确的版本把它当成功能基线再逐步迭代优化每一步都有对比对象不容易迷失。二是理论加速比永远要打折扣。im2col理论上比朴素卷积快但内存扩张可能让带宽成为新瓶颈Winograd理论减少乘法次数但加法开销和数值误差都要算进去稀疏化理论上省了算力但访存不规则和控制开销会抵消一部分收益。我见过太多人拿着理论数字向老板汇报落地时发现根本对不上。做优化的人一定要有“打折预期”的心态。三是工具链和硬件能力决定了优化上限。再好的算法思路也要匹配硬件特性和框架支持。建议动手之前先花时间搞清目标芯片的计算峰值、带宽、缓存大小、SIMD宽度搞清楚推理框架是否支持你想要的融合或稀疏格式。这个准备越充分后面的弯路越少。最后分享一个小技巧做算子优化时记得把“精度差异”也纳入回归测试范围。很多优化会在无意中改变数值计算顺序导致输出结果出现微小偏差。这些偏差在单层测试中看不出来但多层累积就可能让最终精度不可用。所以除了性能跑分一定还要写一个精度对比的自动化用例——输入相同数据优化前后的输出差异要在可接受范围内。这个习惯帮我避免过好几次“性能很好看但根本不能用”的事故。