训练图像模型的时候很多人第一反应是 GPU 不够快、模型不够轻但真正卡住吞吐的往往不是这些。你把nvidia-smi开着盯一会儿就会发现显卡利用率在 60% 到 80% 之间反复横跳而top里面那几个 DataLoader 进程已经把 CPU 的核心吃满了。这个现象在图像分类、检测、分割任务里都极其常见源头就是 JPEG 解码——每张图都要走一遍熵解码、去量化、逆 DCT、色度上采样、颜色空间转换最后才变成 GPU 能用的 RGB 张量。而这张图在几毫秒之后就会被数据增强裁掉一大半、或者被模型压缩成更小的特征图。也就是说大量的解码计算被白白浪费了。JPEG 图像格式加速神经网络训练这个方向做的事情就是利用 JPEG 文件里本来就存着的那份中间产物——去量化后的 DCT 系数跳过逆变换直接把它喂给网络。更激进一点的做法是使用 DCT 训练 CNN把网络的第一层也搬到频域用逐点相乘替代滑窗卷积。这条路不是玄学它建立在离散余弦变换的正交性上数学上是可以严格推导的。下面我把整条链路从原理到代码拆开讲一遍包括我踩过的坑和省不下来的那部分适合对训练加速感兴趣、且手上正好有大量 JPEG 数据要处理的人参考。1. 训练卡在CPU上先把JPEG解码这笔账算清楚1.1 一张图的解码路径到底有多长随便挑一个主流的 C 语言 JPEG 解码库它内部处理一张 224×224 的彩色图大致要走这么几步读文件头、Huffman 熵解码得到量化后的整数系数、乘量化表去量化、对每个 8×8 块做二维逆 DCT、把 Cb/Cr 从 4:2:0 上采样回原分辨率、YCbCr 转 RGB、最后按行写出。整条链路里逆 DCT 和上采样是计算密度最高的两段。我做过粗测单线程下一张 224×224 的 JPEG 完整解码大概落在 0.6 到 1.2 毫秒之间具体取决于压缩质量、色度采样方式和图片内容复杂度。如果训练吞吐是每秒 3000 张那光解码就要吃掉 2 到 4 个 CPU 核心的满负荷这还没算数据增强和内存拷贝。很多团队的做法是开 8 个 worker用核心数硬扛机器上 CPU 占用常年 100%风扇一直转。问题在于这份算力里有一大块是可以用更便宜的方式替代的。逆 DCT 是一组固定的矩阵乘法它在 CPU 上做效率不高但搬到 GPU 上批量做几乎是免费的。色度上采样和颜色转换同理。真正必须留在 CPU 上的只有熵解码和去量化——这两步是变长解码并行度差GPU 不擅长。1.2 JPEG文件里那份没人用的中间产物关键点在这里JPEG 编码的最后一步是熵编码编码的对象是量化后的 DCT 系数。一个 8×8 的块产生 64 个整数系数按 ZigZag 顺序排列大量高频系数被量化成 0最后用 Huffman 表压缩成码流。解码的时候这份系数是先在内存里还原出来的。绝大多数解码接口把它当作内部中间结果直接接着做逆变换用完就扔。但主流解码库基本都提供了另一条路径——直接把这份量化后的系数数组交给你绕过后续所有步骤。拿到的就是一份形状为(块行数, 块列数, 64)的整数数组按通道分开。你不需要逆 DCT、不需要上采样、不需要颜色转换落盘的数据量也小得多。这一步能省掉多少计算我的实测是整条解码链路的时间大概能压掉一半甚至更多剩下的大头是熵解码本身那部分省不了。1.3 收益边界哪些能省哪些省不掉先把话说在前面免得到时候期待落空。这条路线能省的东西是明确的CPU 上的逆 DCT、上采样、颜色转换可以整体搬到 GPU或者干脆不做主机内存到显存的传输量如果传的是系数而不是像素数据量会变化取决于你怎么存下面细说第一层卷积的部分计算如果只保留低频系数计算量能成倍下降。省不掉的东西也明确熵解码变长解码必须逐字节处理很难并行。这是硬成本磁盘 IO文件还是那么大读取时间不变量化带来的信息损失这是有损压缩高频信息已经没了模型只能用剩下的部分。想把这个损失压到可接受范围得控制压缩质量。所以这条路的定位很清楚它不是一个免费加速的魔法而是一次计算位置的重排——把 CPU 上低效、串行、重复的计算挪到 GPU 上高效、批量地做顺便利用频域表示的特性砍掉一部分第一层的计算。想清楚这一点后面所有的设计取舍就都能理解了。2. 正交变换的秘密DCT域里做内积为什么结果不变2.1 二维DCT写成矩阵形式更清楚教科书上的二维 DCT-II 公式看着吓人但写成矩阵乘就非常干净。对一个 8×8 的图像块 $f$定义变换矩阵 $C$$$C[u][x] c_u \cos\frac{(2x1)u\pi}{16}, \quad c_0\sqrt{\tfrac{1}{8}},\ c_{u0}\sqrt{\tfrac{2}{8}}$$那么正向变换就是 $F C f C^{\mathsf{T}}$逆向是 $f C^{\mathsf{T}} F C$。这里有个细节值得单独说一下JPEG 标准里写的系数是 $\frac{1}{4}\alpha(u)\alpha(v)$看着和上面不一样但你把它展开算一遍就会发现8×8 的情况下两者完全等价。也就是说JPEG 用的 DCT 矩阵本身就是正交归一化的$CC^{\mathsf{T}} I$。这一点非常关键后面的等价性推导全靠它。实际代码里把这个矩阵算出来只要几行import numpy as np def dct_matrix(n8): C np.zeros((n, n), dtypenp.float32) for u in range(n): cu np.sqrt(1.0 / n) if u 0 else np.sqrt(2.0 / n) for x in range(n): C[u, x] cu * np.cos((2 * x 1) * u * np.pi / (2 * n)) return C C8 dct_matrix(8) assert np.allclose(C8 C8.T, np.eye(8), atol1e-6) # 正交性自检顺手加一句自检是有必要的因为不同库的缩放约定确实存在差异如果你拿到的系数是某个库自己定义的版本先做一次往返验证正变换再逆变换看能不能还原再往下走。2.2 内积不变性块内卷积可以直接在频域做现在来看核心结论。假设有一个 8×8 的输入块 $f$ 和一个同样 8×8 的卷积核 $w$我们要算的是块内逐元素相乘再求和$$y \sum_{x}\sum_{y} f(x,y), w(x,y) \mathrm{tr}(f^{\mathsf{T}} w)$$把 $f C^{\mathsf{T}}FC$ 和 $w C^{\mathsf{T}}WC$ 代进去其中 $W CwC^{\mathsf{T}}$ 是卷积核的 DCT$$\mathrm{tr}\big((C^{\mathsf{T}}FC)^{\mathsf{T}}(C^{\mathsf{T}}WC)\big) \mathrm{tr}(C^{\mathsf{T}}F^{\mathsf{T}}CC^{\mathsf{T}}WC) \mathrm{tr}(C^{\mathsf{T}}F^{\mathsf{T}}WC)$$再用一次迹的循环性质把末尾的 $C$ 挪到前面和 $C^{\mathsf{T}}$ 抵消$$ \mathrm{tr}(F^{\mathsf{T}}WCC^{\mathsf{T}}) \mathrm{tr}(F^{\mathsf{T}}W) \sum_{u}\sum_{v} F(u,v),W(u,v)$$结论出来了在 DCT 域逐点相乘再求和和在空域逐点相乘再求和结果完全一样。这不是近似是恒等。这就是离散正交变换版本的 Parseval 关系。这条性质怎么用到网络上关键在于卷积的形态。一个kernel_size8, stride8的卷积层输出特征图上每个位置对应的就是一个 8×8 块与卷积核的内积。它恰好落在这个恒等式的适用范围内。而现代视觉模型里patch embedding 这一层通常就是kernel16 stride16或者kernel8 stride8的大步长卷积——就是为了把图像切成不重叠的块。这两件事一拍即合。我一开始也是将信将疑写了个小脚本对比了空域卷积和频域内积的输出最大绝对误差在1e-5量级float32 精度下确认没问题才敢往下做。2.3 低频截断用四分之一的系数换回大部分信息既然内积不变那能不能少用点系数可以但要接受近似。JPEG 量化后能量高度集中在低频左上角那个 DC 系数是块均值的 8 倍携带的能量通常占整块的绝大部分越往右下走系数越小、被量化成 0 的概率越高。这里要注意一个容易搞混的点JPEG 的 ZigZag 顺序不是矩阵顺序。ZigZag 是把 8×8 矩阵按对角线来回扫描排成一维的 64 个位置索引 0 是 DC索引 1 和 2 是最低频的 AC以此类推。你在做低频截断之前必须先把这 64 个值摆回 8×8 的二维布局或者准备好一张索引映射表。截断的收益是可以算的完整的 8×8 内积是每块每通道 64 次乘加只取左上 4×4 就降到 16 次直接省掉四分之三。代价是丢掉高频细节。对分类任务来说这个代价往往可以接受——分类更依赖全局的形状和纹理走向而不是一个像素级的锐利边缘。但对超分辨率、去噪、细粒度识别这类任务截断要慎重我一般建议先做完整的 8×8等确认精度没问题再考虑截断压缩。2.4 算力账面到底省了多少次乘加把账算清楚免得自我感动。假设输入 3 通道 224×224用kernel8 stride8的卷积打 patch输出通道 $C_{out}$环节标准路径DCT 路径完整 8×8DCT 路径截断 4×4CPU 解码到像素需要不需要不需要逆 DCTCPU 上每块 8×8 二维变换省掉省掉块数784784784每块乘加次数64×3×$C_{out}$64×3×$C_{out}$16×3×$C_{out}$是否需要回空域是否否看清楚第三行和第四行完整 8×8 的频域内积乘加次数和空域卷积是一样的。别指望这里能省算力它省的是别的东西——逆变换和上采样。真正能省乘加的是截断4×4 截断能省 75% 的第一层计算量。那这个别的东西值多少我自己的感受是在 CPU 是瓶颈的场景下收益非常实在CPU 侧解码时间砍掉一半GPU 上多做一步矩阵乘几乎感知不到。而在 GPU 是瓶颈的场景下只有截断能带来收益完整 8×8 基本没有加速效果——这一点必须提前想清楚别做完一整套改造发现瓶颈根本不在那边。3. 把解码管线砍掉一半直接读DCT系数3.1 用现成接口拿到量化后的系数主流解码库都提供了跳过逆变换的接口。以常见的 C 语言 JPEG 解码库为例它有一个专门读取系数的函数返回的是一个二维指针数组索引是分量和块行每个元素指向一排 64 个整数系数。用它之后你不需要jpeg_start_decompress那套输出流程直接进入jpeg_read_coefficients的世界。这里有个坑要提前说返回的系数是按分量分开的。Y 分量是每个 8×8 块一份Cb/Cr 在 4:2:0 采样下是每 16×16 像素区域一份因为色度下采样了两个方向。你要自己按分量组织张量不能想当然地当成同构的三通道拼在一起。我第一次写的时候直接把三个分量的块数当成一样结果对不上排查了半天才想起来是色度下采样。另外一个坑是系数的数值范围。这些是整数去量化之前大致落在 [-1024, 1023]去量化之后会更大。存成 int16 一般够用但如果你顺手做了个全局缩放记得把这个缩放记录下来不然统计量全乱。3.2 去量化与块排布的处理细节去量化本身很简单就是每个位置乘上量化表对应位置的值。量化表有两张亮度一张色度一张还存在缩放的版本某些编码器会做质量缩放。处理顺序是从文件头读出量化表注意可能是缩放过的需要按 DQT 段里的精度位判断是 8 位还是 16 位把 ZigZag 顺序映射回 8×8 矩阵顺序逐元素乘量化表得到的就是标准的正交归一化 DCT 系数。这里可以做一个优化如果确定后续只取低频 4×4那就只需要读前若干位 ZigZag 系数后面全是 0 或者不需要。这一步能和熵解码合并因为熵解码本来就是按 ZigZag 顺序产出系数的读到够了就可以提前中断——不过要小心Huffman 码流是连续的提前中断意味着你不能直接跳到下一个块的起始位置需要维护解码状态。如果只是想省事还是老老实实解完 64 个系数。……这样写下去篇幅会失控回到正题。块排布上我建议统一成(C, H/8, W/8, 64)的形状其中最后一维已经是二维布局或者至少你已经能按(u,v)取到。这个形状的好处是后面不管做逐点内积还是接卷积层reshape 都很自然。3.3 在GPU上补完剩下的解码步骤拿到系数之后有两条路。第一条路是保守方案在 GPU 上把逆 DCT、上采样、颜色转换做完还原成普通的 RGB 张量后面的网络一行代码都不用改。我推荐所有人先走这条路因为它风险最低先验证省掉 CPU 这部分到底能快多少再决定要不要往深里做。逆 DCT 在 GPU 上就是两次矩阵乘用einsum写出来非常直观import torch def idct_8x8(coeff, C): # coeff: (B, C, bh, bw, 8, 8) 未归一化的频域系数 # C: (8, 8) 正交 DCT 矩阵 # 返回: (B, C, bh*8, bw*8) b, c, bh, bw, _, _ coeff.shape f torch.einsum(ux,bcuvuv-bcxy, C, x) # 写法示意实际维度需对齐 ...上面这段伪代码里维度写错了正确的应该是把最后两维当成矩阵的两个轴分别乘def idct_8x8(coeff, C): # coeff: (B, C, bh, bw, 8, 8) b, c, bh, bw, _, _ coeff.shape f torch.einsum(ux,bcijxy,vy-bcijuv, C, coeff, C) f f.reshape(b, c, bh * 8, bw * 8) return feinsum会自动把(bh, bw)这两维当作批量维广播不需要手动循环。实测在批大小 64、3 通道、224×224 的设置下这一步在 GPU 上耗时在微秒级基本可以忽略。相比之下在 CPU 上做同样的 784×3 次 8×8 逆变换耗时是毫秒级的——这就是那个数量级的差距来源。至于颜色空间转换YCbCr 到 RGB 就是一个 3×3 的矩阵加偏移写成一次线性变换就行。色度上采样如果是不重叠的 2×2 块复制用repeat_interleave就能完成。要注意上采样的插值方式和解码库原生的方式可能不一致有的是简单复制有的是三角滤波如果对像素级精度敏感这一点要核对。3.4 DCT域里的数据增强怎么写这一步是很多人会卡住的地方。既然不上采样回空域那随机翻转、随机裁剪这些增强怎么办答案是在频域做而且有严格的对应关系。先说翻转。对水平翻转新图像满足 $f(x,y) f(7-x, y)$。代入 DCT 定义$\cos\frac{(2(7-x)1)u\pi}{16} \cos\left(u\pi - \frac{(2x1)u\pi}{16}\right) (-1)^u\cos\frac{(2x1)u\pi}{16}$。所以频率索引 $u$ 的符号会被翻转$$F(u,v) (-1)^u F(u,v)$$同理垂直翻转是 $(-1)^v$两者都做就是 $(-1)^{uv}$。代码上只要提前算好一个 8×8 的符号表做一次逐元素相乘即可比在空域翻转还快sign_h torch.tensor([(-1) ** u for u in range(8)]).view(8, 1).float() sign_v torch.tensor([(-1) ** v for v in range(8)]).view(1, 8).float() def hflip_dct(coeff): # coeff: (..., 8, 8)块的顺序也要整体倒序 return coeff * sign_h def vflip_dct(coeff): return coeff * sign_v注意别漏掉块顺序的倒序——整图翻转时第 j 列块会变成第 (W/8-1-j) 列块系数符号翻转只是块内的部分。再说裁剪。如果裁剪边界对齐到 8 像素那就是纯粹的块索引切片零成本。不对齐的话就得先逆变换再裁或者接受一点近似。实践中我建议把 RandomResizedCrop 的边界对齐到 8损失一点随机性换来完全跳过逆变换很划算。亮度、对比度调整这类颜色抖动在频域里对应的是 DC 系数缩放亮度和 AC 系数整体缩放对比度直接乘系数就完事。Mixup 是线性操作频域插值和空域插值等价也能直接做。4. 网络改造把patch embedding挪到频域4.1 三种接法的取舍改造网络前端有几条路难度和收益差别很大我按实施成本排一下接法改动量是否需要回空域风险适用场景A. GPU 上逆变换后接原网络只改数据层是极低先验证解码加速效果B. 频域内积 stem 原网络改第一层否中patch 对齐 8 的骨干网C. 全频域前端 截断改前几层否较高追求极限吞吐方案 A 我强烈建议所有人都先做一遍它能让你在一两天内看到真实的端到端加速比而不是在理论上打转。方案 B 是真正的DCT 训练 CNN核心就是用频域内积替换第一层的大步长卷积。方案 C 属于优化收益主要在乘加次数上适合已经跑通 B 之后继续压榨。4.2 卷积核变换的推导与实现方案 B 的关键只有一步把空间域的卷积核搬到频域。给定第一层权重w形状(Cout, Cin, 8, 8)对每个(out, in)组合做一次二维 DCT$$W_{freq}[o,i] C , w[o,i] , C^{\mathsf{T}}$$代码用einsum一行搞定def kernel_to_dct(w, C): # w: (Cout, Cin, 8, 8) return torch.einsum(ux,oixy,vy-oiuv, C, w, C)前向传播时输入是一批系数F形状(B, Cin, 8, 8)这里假设整图已经被切成 8×8 块块维度合并进 batch 或者单独保留输出就是$$y[o] \sum_{i}\sum_{u}\sum_{v} F[i,u,v] \cdot W_{freq}[o,i,u,v]$$这个式子长得和一个 1×1 卷积完全一样——只要把(Cin, 8, 8)展平成Cin*64个通道权重展平成(Cout, Cin*64)就是一次标准的线性层。用nn.Linear或者nn.Conv2d(kernel_size1)都能实现走的是高度优化的 GEMM 内核。然后把偏置加上过激活函数后面的层原样接。整个改造的工作量就是重写第一层 改数据加载器剩下的全靠框架自动求导。有个细节需要注意DCT 域的权重也需要初始化策略。因为权重是在频域定义的用空域那套 Kaiming 初始化直接套过来方差会不对劲。我的做法是先按空域规则初始化再做变换或者干脆在频域用正交初始化然后检查一遍输出的激活值分布是否和空域版本接近。这一步不做的话训练初期 loss 会明显比基线高但通常几十个 step 之后就追平了。4.3 归一化的坑DC 和 AC 完全是两种分布这是我在实践中踩得最深的一个坑。频域系数的统计特性和像素完全不是一个量级DC 系数左上角是块均值的 8 倍范围大、均值不为零、分布接近正态AC 系数其余 63 个均值接近 0方差随频率升高迅速衰减高频那一片基本全是被量化到 0 的小值。如果你在数据层做了一次全局的(x - mean) / stdDC 通道会被压得几乎没信号AC 通道被放大成噪声。正确做法是按通道单独统计对每个(c, u, v)组合算各自的均值和标准差64×3 组参数存成一个文件在数据层用。或者更省事直接在网络第一层后面接一个可学习的逐通道仿射本质上就是 BN 去掉批维让网络自己调整。还有一个更朴素的问题预训练权重没法直接用。如果你打算拿一个在大规模数据集上训好的模型做初始化第一层是空域权重你得把它变换到频域才能加载。这一步我在 4.2 的代码里已经写好了反向操作一下就行。但要注意如果原模型第一层的步长不是 8你没法无损转换只能重新训练那一层。4.4 训练超参的调整建议尺度变了学习率和权重衰减都得跟着调。经验上是这样的因为第一层的输入数值范围比像素域大一个量级左右如果沿用原来的学习率初期梯度会比较大容易出现 loss 抖动。我一般会把第一层的学习率单独降到整体的 1/5 到 1/10用参数组分开放其余层保持原设置。等训练稳定下来比如两三个 epoch 之后再放开或者干脆用 warmup 把那一段平滑掉。还有一点是关于BatchNorm的位置。如果你把第一层后面直接接了 BN那频域的尺度问题基本能被它吸收掉归一化的手工调参可以省很多。但要注意 BN 的 running stats 在频域下收敛得会更慢一些因为不同频率通道的动态范围差异大可以考虑把 momentum 调小一点。5. 实测数据与踩坑清单5.1 加速比拆解时间到底花在哪了我在一个中等规模的数据集上和标准管线做了对比硬件是 8 核 CPU 加一张中端卡批大小 64输入 224×224质量因子 90。数据管线部分4 个 worker的耗时大概是这样的阶段标准管线系数管线说明读文件基准基准磁盘 IO 不变熵解码 去量化基准基准变长解码省不掉逆 DCT 上采样 转换基准0搬到 GPU主要收益点主机到显存传输像素张量系数张量视存储格式而定GPU 侧逆变换0很小微秒级端到端的表现是在 CPU 本来就是瓶颈、GPU 利用率只有 70% 左右的场景下改造后 GPU 利用率能拉到 90% 以上整体吞吐有显著提升。而在原本 GPU 就吃满的场景下几乎没有变化——这点前面已经讲过这里再强调一次因为很多人做完改造发现没加速就觉得是方法有问题其实就是瓶颈判断错了。判断方法很简单把批大小调到很小、把数据增强全部关掉看 GPU 利用率。如果已经接近 100%那瓶颈在 GPU别折腾数据管线了。5.2 精度掉点与恢复手段用质量因子 90 的数据、完整 8×8 系数、不做截断我这边复现的结果是和像素域基线基本持平差异在随机种子的波动范围内。但当质量因子降到 75 以下精度会出现肉眼可见的下滑尤其是那些依赖细节纹理的类别。如果确实要低质量数据有几个补救方向一是不要用截断完整保留 64 个系数二是把高频那些被量化成 0 的通道显式作为一个 mask 输入让网络知道哪些位置的信息不可靠这个技巧对小目标检测特别有效三是用渐进式加载先加载高质量子集训练再混入低质量数据微调。第三条是我目前觉得最省事的。另外色度下采样也是一块隐形的信息损失。4:2:0 意味着色度分辨率只有一半如果你的任务对颜色敏感比如某种特定的材质区分要考虑换成 4:4:4 采样的数据源或者干脆在频域里把两个色度块直接合并到 Y 的块结构里。5.3 常见问题速查现象大概率原因处理方式前向输出和空域版本对不上DCT 矩阵缩放约定不一致先做往返自检确认 $CC^{\mathsf{T}}I$训练 loss 一开始就爆炸频域输入尺度没归一化按(c,u,v)分别做标准化精度比基线低 5 个点以上压缩质量太低或做了截断换高质量数据取消截断加速比不明显瓶颈不在 CPU关掉增强测 GPU 利用率色度块数对不上4:2:0 下采样被忽略了按分量分开组织张量翻转增强后图像位置错乱只翻了块内系数没翻块顺序系数符号变换 块索引倒序加载预训练权重报错第一层形状不匹配用权重变换函数转换后再加载显存占用反而变大系数用 float32 存、通道数×64存成 int16 或做稀疏压缩最后一行值得多说一句。系数张量的形状是(B, 3, 28, 28, 64)看起来比(B, 3, 224, 224)的元素多其实是完全一样的28×28×64 224×224只是排布不同。但如果你按 float32 存、又给每个频率通道单独做了归一化中间缓存会膨胀。解决办法是在数据层保持 int16进网络之前再转 float 并做缩放这一步的转换开销很小。还有一个不那么常见但会浪费很多时间的问题部分编码器会在文件里写入非标准的量化表或重启标记间隔。如果你的解码路径在处理重启标记时状态没重置好系数会出现周期性错位表现为图像上每隔一段就有一条乱块。这种情况在批量处理各种来源的数据时偶尔会遇到排查的方法是先把出问题的文件单独解一遍和标准的完整解码做像素级对比定位到具体是哪一段开始出错的。我个人在整个过程中的体会是这条路线的技术门槛不在数学上——正交变换保持内积这个结论写过一遍推导就永远不会忘——而在工程细节上。块排布、色度下采样、量化表版本、归一化统计、翻转时的块顺序每一个都是不起眼但能卡住你半天的小问题。我的建议是先把方案 A 跑通用一个最简的对比脚本验证系数 → 逆变换 → 像素和直接解码 → 像素完全一致把这个地基打牢再往上叠频域的第一层。地基不稳的时候去调网络你会怀疑人生因为你根本分不清是数学错了还是数据错了。至于那个低频截断的开关先放着等精度确认没问题了再说——很多时候完整 8×8 已经够用没必要为了省那点乘加去赌精度。
