两个月前,我接到一个画质增强服务的模型选型任务。需求简单粗暴:把经典超分模型的 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 图像就是 255,MSE 是重建图和 GT 之间的像素均方误差。也就是说,PSNR 本质上衡量的是重建结果和原始高分辨率图像之间的像素级误差能量。
反推一下就能知道 0.17dB 到底意味着什么。假设某个测试集上原模型 PSNR 是 28.00dB,涨 0.17dB 后变成 28.17dB,MSE 下降的比例大约是: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的难度 |
|---|---|---|---|
| Set5 | 5 | 人像、动物、简单背景 | 很难,分数接近饱和 |
| Set14 | 14 | 自然图像混合 | 中等 |
| BSD100 | 100 | 日常生活场景 | 中等 |
| Urban100 | 100 | 城市建筑、重复结构 | 相对容易,但模型需要真理解 |
| Manga109 | 109 | 漫画线稿、大色块 | 纹理多,容易涨但过拟合风险也大 |
按 SwinIR 论文公开的参考值,它在 ×4 任务上 Set5 约 32.92dB、Urban100 约 27.45dB、Manga109 约 30.92dB。我们的模型在 Urban100 和 Manga109 上平均提升 0.17dB,Set5 和 Set14 基本持平,这个分布是合理的——简单数据集已经没有空间,复杂结构数据集还能榨出油水。
1.3 掉点红线:小模型的否决条件
轻量化实验里有个不成文的规矩:减了参数可以接受小幅掉点,但几个关键数据集不能同时崩。我们实验过程中定的红线是:Set5 掉点不得超过 0.05dB,BSD100 掉点不得超过 0.03dB,否则方案直接否掉。最终结果是在 Set5 上基本持平,BSD100 微涨,完全踩线通过。这个红线不是拍脑袋定的,因为产品侧的体感阈值就在这个范围附近——大模型和小模型的重建结果差异如果低于 0.1dB,用户基本看不出区别,但参数减少 12% 带来的显存和延迟收益却是实打实的。
2. SwinIR的参数量都花在哪了:减12%要先知道肥肉在哪
2.1 从SwinIR的骨架说起
SwinIR 的结构分成三段:浅层特征提取、深层特征提取、重建模块。浅层特征提取是一个 5×5 卷积,把输入图映射到高维特征空间;深层特征提取由 6 个 RSTB(Residual 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_dim=180,MLP ratio=2,window_size=8。单层参数构成大致如下:
| 组件 | 参数计算 | 单层参数量 | 36层合计 | 占总参比例 |
|---|---|---|---|---|
| QKV投影 | 180×540 + bias | 97.2k | 3.50M | 约30% |
| 输出投影 | 180×180 + bias | 32.4k | 1.17M | 约10% |
| MLP第一层 | 180×360 + bias | 64.8k | 2.33M | 约20% |
| MLP第二层 | 360×180 + bias | 64.8k | 2.33M | 约20% |
| 两个LayerNorm | 约360+360 | 约0.7k | 约25k | <1% |
| RSTB尾部卷积与Patch Embed | 每个约330k | 共约2M | 约17% |
单层合计约 260k,36 层就是 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 ratio=2,hidden 宽度只有 360,并不算宽,但实际训练中 36 层的 FFN 存在大量相似的线性映射模式——后层经常在前层的基础上做微调。直接降宽度会损失表达能力,但如果配合蒸馏和初始化补偿,这部分潜力和风险并存,值得动刀。LayerNorm 参数很少,动了也没意义,不如留着。
3. 让模型小12%还不掉点的三个做法:共享权重的艺术
3.1 相邻RSTB同层共享QK投影:节省约10%参数
这是整个压缩方案里最核心的一刀。我们的做法是:把 6 个 RSTB 按顺序组成 3 对,每一对 RSTB 的同一深度层共享同一个 QK 投影矩阵,V 投影和输出投影保持独立。SwinIR 的窗口大小在所有层都是 8,embed_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 降到 330,36 层合计节省约 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_feat=True) 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% 的显存和延迟收益值钱,做之前先想清楚你要交付的到底是论文里的表格,还是一个能扛住线上压力的服务。