两个月前我接到一个画质增强服务的模型选型任务。需求简单粗暴把经典超分模型的 PSNR 再往上顶一点同时把显存占用压下来。第一反应肯定是在 SwinIR 这种 Transformer 架构上做文章——过去两年图像超分领域基本被它和它的变体统治精度确实高但一个标准模型 11.8M 参数在视频类业务里跑一圈显存和时间成本都不太好看。后来我们在 SwinIR 上做了一轮结构压缩加蒸馏实验最后交付的模型比原版小 12%在 Urban100 和 Manga109 这类结构纹理密集的数据集上平均 PSNR 还反超 0.17dB。这篇文章就把整套做法、训练技巧和踩坑记录完整捋一遍适合正在做超分模型轻量化、或者想在 Transformer 超分基础上继续涨点的同学参考。1. 0.17dB的含金量先搞清楚这条涨点值不值钱1.1 PSNR这个数字是怎么算出来的超分领域聊涨点几乎都是拿 PSNR峰值信噪比说话。公式不复杂PSNR 10 × log10(MAX² / MSE)其中 MAX 是像素最大值比如 8bit 图像就是 255MSE 是重建图和 GT 之间的像素均方误差。也就是说PSNR 本质上衡量的是重建结果和原始高分辨率图像之间的像素级误差能量。反推一下就能知道 0.17dB 到底意味着什么。假设某个测试集上原模型 PSNR 是 28.00dB涨 0.17dB 后变成 28.17dBMSE 下降的比例大约是1 - 10^(-0.17/10) ≈ 3.8%。换句话说整张图的像素误差能量被压低了接近 4%。放在一张 1000 万像素的照片里这个收益是肉眼能在边缘和纹理区域感受到差异的幅度但又不会像 0.5dB 那样一眼惊艳。学术界和工业界的通用判断是0.1dB 算有效涨幅0.3dB 是可以写进版本发布说明的进步0.5dB 以上基本意味着架构级别的突破。所以 0.17dB 这个数字放在模型还变小了的背景下含金量并不低。1.2 不同测试集上的0.17dB含金量完全不同同样是涨 0.17dB在不同测试集上的难度和说服力差别很大。超分领域有几个被反复刷的老测试集Set5、Set14、BSD100、Urban100、Manga109。Set5 只有 5 张图内容简单模型分数早就被刷到接近饱和想再涨 0.17dB 非常困难Urban100 是 100 张城市建筑图充满了重复结构、线条和窗户网格模型必须真正理解结构规律才能涨分Manga109 是 109 张漫画有大面积均匀色块和密集线稿同样属于吃结构理解的数据集。测试集图像数量内容特点涨0.17dB的难度Set55人像、动物、简单背景很难分数接近饱和Set1414自然图像混合中等BSD100100日常生活场景中等Urban100100城市建筑、重复结构相对容易但模型需要真理解Manga109109漫画线稿、大色块纹理多容易涨但过拟合风险也大按 SwinIR 论文公开的参考值它在 ×4 任务上 Set5 约 32.92dB、Urban100 约 27.45dB、Manga109 约 30.92dB。我们的模型在 Urban100 和 Manga109 上平均提升 0.17dBSet5 和 Set14 基本持平这个分布是合理的——简单数据集已经没有空间复杂结构数据集还能榨出油水。1.3 掉点红线小模型的否决条件轻量化实验里有个不成文的规矩减了参数可以接受小幅掉点但几个关键数据集不能同时崩。我们实验过程中定的红线是Set5 掉点不得超过 0.05dBBSD100 掉点不得超过 0.03dB否则方案直接否掉。最终结果是在 Set5 上基本持平BSD100 微涨完全踩线通过。这个红线不是拍脑袋定的因为产品侧的体感阈值就在这个范围附近——大模型和小模型的重建结果差异如果低于 0.1dB用户基本看不出区别但参数减少 12% 带来的显存和延迟收益却是实打实的。2. SwinIR的参数量都花在哪了减12%要先知道肥肉在哪2.1 从SwinIR的骨架说起SwinIR 的结构分成三段浅层特征提取、深层特征提取、重建模块。浅层特征提取是一个 5×5 卷积把输入图映射到高维特征空间深层特征提取由 6 个 RSTBResidual Swin Transformer Block串成每个 RSTB 内部又包含一个 Patch Embed、6 层 Swin Transformer Layer 和一个残差卷积最后重建模块用卷积加 pixel shuffle 把特征升到目标分辨率。整套下来标准的 SwinIR 参数量约 11.8M。其中绝大部分参数不在浅层卷积也不在重建模块而在那 36 层 Swin Transformer Layer 里。如果你要压缩一个 Transformer 超分模型首先得搞清楚每层参数都花到了哪里而不是盲目地砍层数或者缩通道数——层数一砍感受野和建模深度立刻受损通道数一缩每层的特征表达能力整体下降。精准压缩应该找到冗余最大、影响最小的部分。2.2 每个Transformer层的参数构成一个标准的 Swin Transformer Layer 包含窗口自注意力W-MSA/SW-MSA和两层 MLP。以 SwinIR 经典配置为例embed_dim180MLP ratio2window_size8。单层参数构成大致如下组件参数计算单层参数量36层合计占总参比例QKV投影180×540 bias97.2k3.50M约30%输出投影180×180 bias32.4k1.17M约10%MLP第一层180×360 bias64.8k2.33M约20%MLP第二层360×180 bias64.8k2.33M约20%两个LayerNorm约360360约0.7k约25k1%RSTB尾部卷积与Patch Embed每个约330k共约2M约17%单层合计约 260k36 层就是 9.36M加上 RSTB 内部的卷积和 Patch Embed 以及首尾模块刚好拼出 11.8M 的总数。从表格可以看得很清楚注意力部分QKV 输出投影占 40%FFN 占 40%剩下 20% 是卷积和嵌入。所以压缩的核心战场就是注意力和 MLP 的投影矩阵想减 12% 参数还不伤精度必须在这两个地方找冗余。2.3 哪些参数是可以动的哪些不能动我的经验是输出投影和 V 投影基本不能乱砍它们直接影响恢复出来的特征质量Q 和 K 投影反而是可以动手的地方。原因在于Q 和 K 的作用是计算 token 之间的相关性分数它们学习到的是哪些位置应该互相参考的通用关系这种关系在不同层之间高度相似。真正让层与层之间产生差异的主要是 V 投影和 FFN。另外FFN 是参数量最大的部分但也是冗余最明显的部分。SwinIR 里的 MLP ratio2hidden 宽度只有 360并不算宽但实际训练中 36 层的 FFN 存在大量相似的线性映射模式——后层经常在前层的基础上做微调。直接降宽度会损失表达能力但如果配合蒸馏和初始化补偿这部分潜力和风险并存值得动刀。LayerNorm 参数很少动了也没意义不如留着。3. 让模型小12%还不掉点的三个做法共享权重的艺术3.1 相邻RSTB同层共享QK投影节省约10%参数这是整个压缩方案里最核心的一刀。我们的做法是把 6 个 RSTB 按顺序组成 3 对每一对 RSTB 的同一深度层共享同一个 QK 投影矩阵V 投影和输出投影保持独立。SwinIR 的窗口大小在所有层都是 8embed_dim 也完全一致因此矩阵形状天然匹配不需要任何额外对齐操作。为什么敢这么干训练完成后的 SwinIR不同 RSTB 的同层特征在激活分布上有很强的相似性尤其是相邻 RSTB——它们处理的是同一分辨率下的特征图分工接近只是不断在细化。Q 和 K 投影学习的是 token 间的相关模式这种模式本质上和当前特征处于哪一层关系不大更像是通用结构先验。共享 QK 相当于给相邻 RSTB 加了一个归纳偏置你们对相关性的理解是一致的差异集中在内容变换上有 V 和 FFN 去表达。参数账很好算。每层原本有 Q 和 K 两个 180×180 矩阵两个层就是四个共享后变成两个每对层节省约 64.8k 参数。36 层配对成 18 对一共节省约 1.17M占总参数近 10%。这一步做了之后模型精度下降很小Set5 上几乎持平Urban100 上大约掉 0.05dB 以内可以说是性价比最高的一刀。3.2 FFN降宽加初始化补偿再省约3%第二刀动 FFN。把 MLP ratio 从 2.0 降到 1.83即 hidden 宽度从 360 降到 33036 层合计节省约 390k 参数占比约 3.3%。两刀加起来大约 13%为了让参数减少的百分比更保守更可信共享 QK 后额外保留一层可训练的偏置项最终对外宣称的数字就是 12%。直接降宽会让学生模型的收敛速度变慢因为每层输出的方差会随着 hidden 宽度缩小而略微下降残差叠加后这个差距会被逐层放大。解决办法很朴素在初始化时把 FFN 第二层down projection的权重乘以 sqrt(360/330) ≈ 1.044让每层输出方差与原始模型保持一致。不要小看这个细节我试过不补偿直接训练前 20k 步 loss 下降明显更慢最终收敛 PSNR 低了约 0.08dB补偿之后基本追平。这里想强调一个原则压缩 Transformer 时精度损失往往不是单一因素造成的而是多个被忽略的小改动叠加的结果。初始化方差匹配是其中一个非常容易被漏掉的点。类似的还有 attention 里的温度缩放、残差 dropout 概率任何改变特征统计分布的动作都要在初始化阶段做反向补偿。3.3 训练时师生蒸馏推理时零成本涨点结构压缩做完之后模型确实小了 12%但精度并没有超过原版。真正把那 0.17dB 净增益拉出来的是蒸馏策略。我们用原始 SwinIR 当教师模型压缩后的 SwinIR-T 当学生模型在训练后半段同时优化两个目标重建损失和特征蒸馏损失。蒸馏的具体做法是教师模型完全冻结只用来提供中间特征。我们取教师倒数第二个 RSTB 的输出作为软标签让学生模型对应位置的 RSTB 输出向它对齐。损失设计很简单L L1(pred, GT) λ × L1(feat_stu, feat_tea)。前期只跑第一项让学生的结构先稳定下来训练进度到 50% 之后开始加蒸馏项λ 取 1.0。# 伪代码蒸馏训练主循环 teacher.eval() for step, (lr_patch, hr_patch) in enumerate(train_loader): with torch.no_grad(): tea_feat teacher.extract_feat(lr_patch, target_depth-2) pred, stu_feat student(lr_patch, return_featTrue) loss_recon F.l1_loss(pred, hr_patch) loss_distill F.l1_loss(stu_feat, tea_feat) loss loss_recon lambda_distill * loss_distill loss.backward()为什么蒸馏能反超因为教师模型在训练过程中已经见过大量数据它的中间特征相当于一个被验证过的优化路径。学生模型参数更少如果只看最后的重建损失自己摸索很容易陷入次优局部解但跟着教师的特征走相当于沿着一条成熟的路快速收敛省下来的参数反而被用在了更关键的位置上。我们实验里蒸馏带来的增益大约有 0.2dB 以上不仅填平了压缩引入的掉点还倒赚了 0.17dB。3.4 为什么这套减法能反超过参数化与归纳偏置的平衡很多人问过我为什么砍了参数反而能涨点是不是实验有 bug。我的理解是SwinIR 本身有过参数化成分36 层 Transformer 在训练数据上并没有完全用完它的表达能力很多参数在学重复的模式。共享 QK 本质上是把这些重复模式参数化地固定下来缩小了搜索空间的同时也减小了过拟合风险。在 Urban100 和 Manga109 这种结构复杂的数据集上过拟合是涨点的大敌归纳偏置反而帮了忙。但这里有个边界条件这种减法只适合在训练数据量中等、模型本身偏大的场景。如果你的模型已经很小比如 SwinIR-light 那种不到 1M 参数的状态再砍参数就是净亏。另外蒸馏对反超的作用也不能夸大它依赖于教师模型的质量教师越强学生的上限越高。我们后面也试过用大模型当教师、小 CNN 当学生的搭配发现收益同样可观这和 Transformer 内部蒸馏是两个方向。4. 训练配方比结构更关键蒸馏、数据增强与实测复现4.1 基础配置数据、batch、优化器很多超分实验跑不出论文数字问题不在模型结构而在训练配置。这里给出一份我们最终使用的配方照着跑基本能复现出小12%还涨0.17dB的结果。训练数据用 DIV2K 加 Flickr2K大约 8000 张高分辨率图随机裁剪 64×64 的 patch批大小 32。优化器用 AdamW初始学习率 2e-4用余弦退火调度降到 1e-5总迭代数 500k。EMA 指数滑动平均设 0.999混合精度训练开启。单卡 A100 上跑一轮实验大概需要 36 到 48 小时如果资源紧张可以把迭代数降到 300k但涨点幅度会缩水大约 0.05dB 到 0.08dB。数据增强方面超分领域常用的就是随机旋转 90/180/270 度和水平翻转我们没加更激进的颜色增强因为超分任务要求模型忠实还原原始颜色过度颜色扰动反而有害。这个配置和官方 SwinIR 基本一致方便做公平对比。4.2 结构重参数化让推理模型更紧凑除了参数数量减少我们还在推理阶段做了一个结构重参数化的小优化。SwinIR 每个 RSTB 尾部有一个 3×3 卷积作为残差连接旁边还并联着主路径的输出。训练时两条路径独立数学上等价于推理时把这两个卷积融合成一个同尺寸卷积。这样的好处是推理图里的算子更少实际跑起来还能再快一点。这部分优化不影响参数数量统计但确实影响部署体验。我们在同样的推理框架下对比过融合后的模型在 1080p 图像上的单帧推理时间比未融合版本减少了约 7%显存峰值低了约 200MB。注意结构重参数化需要在训练结束后重新导出权重且要和 batch norm 这种依赖数据统计的操作一起处理——SwinIR 里没有 BN所以这里很干净直接做卷积加法就行。4.3 验证协议怎么测出可信的0.17dB测试协议照着 SwinIR 论文来Set5、Set14、BSD100、Urban100、Manga109 五个数据集评估 ×2/×3/×4 三档缩放。我们公布的 0.17dB 指的是 ×4 任务在 Urban100 和 Manga109 上的平均涨幅其他数据集基本持平。每次评测用 RGB 空间、不裁剪 border、保留所有像素计算 PSNR这些细节不统一的话数字之间完全没法比。还有一个容易被忽略的点测试时不能开 drop_path。很多训练脚本默认开 drop_path如果测试忘记关掉PSNR 会莫名其妙掉 0.1dB 以上而且每次跑还不一样。检查方法很粗暴对同一张图连续跑两次结果不稳定就说明测试模式有问题。4.4 踩坑清单FP16、蒸馏超参、梯度检查这次实验我们踩了不少坑挑三个最值得说的。第一个是 FP16 下的蒸馏 loss 溢出问题。蒸馏 loss 的量级通常比重建 loss 小一个数量级混合精度训练时容易出现 underflow表现是 loss 曲线前期正常、后期突然出现几次 NaN。解决办法是把蒸馏 loss 乘以一个固定缩放系数后再相加或者在混合精度配置里单独提高该 loss 的动态 scale 上限我们最终选的是后者。第二个是蒸馏权重 λ 不能贪。试过 λ2.0训练后期学生模型完全被教师牵着走失去了自主探索能力结果反而比 λ1.0 低了约 0.07dB。我的经验是 λ 在 0.5 到 1.0 之间最稳。第三个是共享 QK 投影后必须检查梯度。共享参数会同时接收两个层的梯度如果两层对同一个参数的梯度方向长期相反训练会非常慢。我们通过在训练早期打印共享参数的梯度范数来排查发现 RSTB 对内的层梯度方向一致性很高基本不会互相打架。这个验证步骤花不了十分钟但能避免你跑完一轮才发现不收敛。5. 把同样的思路搬去视频超分时间维度是更大的素材库5.1 视频超分为什么更吃显存视频超分正在成为比图像超分更热的方向核心原因很简单视频是连续的图像序列信息量更大训练和推理的代价也更高。典型的视频超分模型如 BasicVSR、BasicVSR、VRT都要同时处理空间和时间两个维度。空间维度的特征提取往往直接复用 SwinIR 里的 Transformer block时间维度则用光流或可变形卷积做帧间对齐。显存压力主要来自两部分一是多帧同时进入网络激活值成倍增加二是时间对齐模块需要保存中间光流和包络特征。SwinIR 作为空间骨干的时候它的尺寸会直接乘以帧数。所以 小 12% 的收益在视频任务里会被放大——如果一次处理 6 帧空间骨干小了 12%整体显存大概能省 8% 到 10%帧率也会有可感知的提升。5.2 把轻量化设计迁移到视频场景我们在视频超分实验里把共享 QK 投影的方案直接移植到 VRT 的空间特征提取部分效果依然成立。VRT 里同样有多个 Transformer block不同 block 处理的是不同时刻或不同分辨率的特征相邻 block 之间的 QK 相关模式同样高度相似。不过要注意一个差异视频模型里有的 block 还承担跨帧注意力这类 block 的 QK 投影最好不要和纯空间 block 共享否则跨帧相关性会被稀释。更值得做的是帧间蒸馏。图像蒸馏是让学生模型贴近教师模型的中间特征视频蒸馏则可以更进一步利用视频的时间一致性作为额外监督。做法是让学生模型重建相邻两帧时输出特征之间的差值也要和教师模型对应帧的特征差值保持一致。这个约束叫 temporal feature consistency loss能有效抑制视频超分常见的闪烁现象。我们做过一个对比实验同样的视频超分小模型只做空间蒸馏时连续帧的 PSNR 波动在 ±0.3dB 左右加上时间一致性约束之后波动降到 ±0.15dB 以内主观上基本不闪了。这个收益在纯图像超分里是看不到的但视频场景里非常关键。5.3 实际部署收益推理帧率提升约18%最后说一个部署侧的实测数字。把压缩后的图像超分模型直接用在视频逐帧超分上对比原版 SwinIR在同样的 1080p 输入下推理帧率提升约 18%显存占用下降约 12%。如果再加上光流缓存和特征复用也就是相邻帧的变化区域复用上一帧的中间特征还可以再省掉约 25% 的重复计算。这里有个取舍值得注意帧间特征复用砍掉的是重复计算但会引入轻微的质量依赖万一某一帧光流估计不准错误会往后续帧传播。我们实际的解决方式是一个置信度判断——光流置信度低的时候强制重新计算该区域特征置信度高的时候才走缓存分支。这个策略让帧率提升没有牺牲边缘帧的稳定性。最后分享一点个人体会做这种轻量化项目我最大的感受是别把小 12%和涨 0.17dB当成两个独立指标它们是同一套方案的上下游。单纯压缩参数很容易但压缩之后还能靠蒸馏和训练策略反超就需要把结构和训练放到一起设计。如果你也想复现这条路我建议首先把官方 SwinIR 仓库跑通拿到一个能复现论文指标的 checkpoint 再动手砍结构——蒸馏教师的质量直接决定学生的上限教师本身不达标后面全白搭。另外0.17dB 在学术视角是有效涨幅在产品视角可能不如那 12% 的显存和延迟收益值钱做之前先想清楚你要交付的到底是论文里的表格还是一个能扛住线上压力的服务。