SwinIR这个基线我用了很久。做图像超分的时候它几乎是默认的Transformer骨架做视频超分的时候也有很多模型拿它当特征提取主干。最近接了一个轻量化超分项目目标是让模型在部署侧跑得动、效果还得保得住。最后交付的模型参数量比SwinIR小了12%在测试集上的PSNR不但没掉反而还高了0.17dB。这篇文章就把整个优化过程复盘一遍为什么拿SwinIR开刀“小12%”这个数字究竟怎么定义出来的又是怎么一步步把模型做小、把指标做高的。这篇内容适合两类人看。一类是做超分模型部署的工程人核心诉求是模型压小之后精度不要回退可以重点关注剪枝和蒸馏的组合方式另一类是准备在Transformer类模型上做压缩的研究者可以参考结构化剪枝、特征蒸馏、重参数化这三板斧怎么搭配。读完你可以拿到一条可复现的完整路径也能避开我在调参时踩进去的那些坑。1. 给超分模型瘦身为什么拿SwinIR开刀1.1 超分任务到底在解决什么问题单张图像超分本质上是一个病态的逆问题。一张低分辨率图可以对应无数张可能的高分辨率图模型要学的不是简单插值而是从大量样本里归纳出“哪种高频细节最合理”——比如皮肤纹理、布料网格、远处文字边缘。早期用SRCNN、VDSR这类全卷积网络效果好但放大倍数一大就发糊后来EDSR、RCAN把残差连接和注意力机制做到极致再往后SwinIR将窗口自注意力引入超分公开基准上的PSNR又上了一个台阶。这个任务的难点在于高频细节是“猜”出来的。低分辨率图像里丢失的信息不会平白无故回来网络只能通过训练集里统计到的先验去补。这也是为什么超分模型普遍偏大、偏深——容量越大能记住的纹理模式越多。但部署场景又是另一回事手机端、边缘盒子、视频后处理流水线都不允许模型动辄几十兆参数。所以超分模型的“压缩”不是一个可选项而是落地第一步。1.2 SwinIR为什么能当这个基线SwinIR的结构可以拆成三大段。浅层特征提取是一个3x3卷积深层特征提取由多个RSTBResidual Swin Transformer Block堆叠而成每个RSTB内部包含若干个Swin Transformer Block还带残差连接和卷基层最后是重建模块通过PixelShuffle把特征上采样回目标尺寸。这套设计的关键就两点一是窗口自注意力把全局建模成本控制在窗口大小内分辨率增大时不会计算爆炸二是移动窗口机制让相邻窗口之间能交换信息图像这种局部相关性很强的数据特别吃这一套。为什么选它当基线我做技术选型时只看三件事。第一效果是否公认。SwinIR在超分领域被引用次数摆在那拿它作为“比SwinIR更强”的参照物审稿人、同事、领导都认。第二代码是否成熟。官方的PyTorch实现干净没有历史包袱我改起来不需要花一周时间逆向别人的魔改代码。第三冗余度够不够高。RSTB堆叠数量一多参数量就集中在少数几个线性层和卷积上压缩空间很大。选一个“既强又有水分”的基线后面做减法才有意义选一个本身就轻量的小模型压无可压折腾半天只省下几KB没有说服力。1.3 压缩目标不是简单把网络变薄很多人一上来就把所有卷积的输出通道减半以为模型小了就完事。真实情况远没那么简单。Transformer类模型里注意力头数量、窗口大小、FFN中间层维度、卷积通道数这些维度相互耦合。你动其中一个其他部分的形状都得跟着变否则前一层输出通道数和后一层输入通道数对不上整个网络就直接跑不起来。我这次项目定的方向很明确保持SwinIR的整体宏架构不动在RSTB内部做通道级别的结构化裁剪再用蒸馏把精度拉回来。宏架构保留意味着我还可以复用SwinIR的预训练权重、官方训练策略、甚至部分中间特征做对齐内部通道裁剪则保证了参数和计算量实打实往下降。这个思路适合绝大多数“以SwinIR为baseline”的项目不需要重新发明轮子。2. “小12%”和“涨0.17dB”到底是怎么定义的2.1 “小12%”是参数、FLOPs还是模型文件体积做模型压缩最怕的就是指标定义含糊。有人说的“小”是参数量少了有人说的是FLOPs降了还有人说的是模型文件从XX MB变成XX MB这三个概念经常被混着讲实际差别很大。我这次项目里硬性指标是“参数量下降12%”。SwinIR的参数量在11.8M左右不同实现会有少许出入12%意味着最终模型要压到10.4M附近。参数量下降模型文件体积基本跟着等比下降因为权重的字节数就是参数个数乘以每个参数的存储位数。FLOPs我也一起统计了因为通道剪掉之后计算量通常会跟着降实测大约下降了14%。至于单张图片的推理耗时我没有把它当作第一版的目标——它受内存带宽、算子调度、推理框架影响非常大不是严格跟参数和FLOPs呈正比的。为什么要先盯参数量因为我当时的部署环境对模型体积和参数读取成本更敏感模型需要做热更新体积小意味着带宽压力小。如果你的瓶颈是GPU算力那应该优先盯FLOPs如果瓶颈是端侧缓存带宽优先盯参数量。先想清楚要优化哪个再决定用什么压缩手段这一步想不清楚后面所有实验都可能白做。2.2 0.17dB的提升在超分领域是什么概念PSNR是基于MSE算出来的对数指标数值越高代表像素重建误差越小。超分领域大家卷来卷去也就是0.1到0.2dB这个区间。SwinIR当年相对RCAN的提升大概是0.3dB级别所以0.17dB已经不是“小数点后的运气”而是一个相当可观的增益尤其在模型还变小12%的情况下出现含金量更高。但我也要泼一盆冷水。PSNR涨0.17dB不代表人眼一定能看出明显差异。超分领域很常见的情况是PSNR上去、主观效果反而变差比如纹理过度平滑、边缘出现振铃。所以这次项目里我不光看PSNR还用局部放大图对比了边缘和文字区域。最终得出结论增益主要来自高频细节更干净尤其栅栏、布料这类重复纹理区域振铃和伪影都比SwinIR少。只看整体PSNR容易骗自己配合主观对比才算数。2.3 压缩方案怎么选型为什么是剪枝蒸馏重参数化能缩小模型的路子很多剪枝、量化、低秩分解、NAS、蒸馏还有重参数化。这些手段不互斥但乱组合会带来巨大的排查成本。比如你先量化再剪枝效果掉得都分不清是谁的锅。我最终定的组合是“结构化剪枝 知识蒸馏 重参数化”原因有清晰的逻辑链。结构化剪枝是直接从架构上删掉冗余通道和注意力头模型是真的“变瘦”不是靠量化去压缩数值精度蒸馏是让剪枝后的学生模型跟着原始大模型学缩小容量差距甚至在某些局部反超重参数化负责解决训练和推理不对等的问题训练时多用几个分支更好收敛推理时合并成单卷积计算量还能再降一截。量化我在第一版没有碰。端侧INT8量化还要考虑校准集、动态范围、反量化开销一步到位很容易收获各种掉点。先把模型结构整理干净后面再加量化才是稳妥顺序。3. 三步走把模型做小的同时把精度做高3.1 第一步结构化剪枝先找出模型里的冗余通道剪枝前必须先做侦测不能闭着眼砍。我做的第一件事是给SwinIR内部所有卷积层计算每个输出通道的L1范数也就是把每个输出通道对应的滤波器权重绝对值求和。权重整体接近0的通道说明学到的特征可有可无优先剪掉权重数值大的通道保留。对注意力头我用的判断标准是该头输出特征的平均绝对值可以理解成每个头的“激活强度”。一个很关键的操作是剪枝不是逐层均匀剪而是全局排序后按比例剪。不同层的冗余程度差别很大有的层冗余通道占30%有的层只占5%。逐层均匀剪会让冗余少的层也被强行削弱白白掉精度。我给每个通道算完分数后把全模型所有分数拼在一起取阈值低于阈值的通道才被剪掉。def compute_channel_scores(model): scores {} for name, module in model.named_modules(): if isinstance(module, nn.Conv2d): # 对该卷积的每个输出通道取权重L1范数 scores[name] module.weight.data.abs().sum(dim(1, 2, 3)) return scores def make_pruning_mask(scores, ratio0.12): all_values torch.cat([s.flatten() for s in scores.values()]) threshold torch.quantile(all_values, ratio) masks {k: (v threshold) for k, v in scores.items()} return masks剪完不是直接把通道删掉而是要生成mask或者构造新的权重矩阵确保前一层输出通道数和后一层输入通道数对得上。SwinIR里的残差连接会把某些通道直接“加回来”如果这些分支上的通道被剪掉剪枝后的结构会跟残差语义冲突特征流就算错了。所以我对参与残差连接的卷积层专门设置了一份“保护名单”这些层只统计分数、不参与剪枝。这一步看起来不起眼却是掉点最常发生的原因千万不能偷懒。3.2 第二步知识蒸馏让大模型带着小模型重新学剪完结构后学生模型的容量已经变小了直接从头训练到收敛通常只能恢复到基线附近想超过老师很难。这是容量限制决定的硬训也突破不了。所以我引入了蒸馏。具体做法是教师用原始SwinIR权重完全固定学生用剪枝后的模型损失函数拆成三块。L_total L_recon α * L_feat β * L_attnL_recon是像素空间重建损失这里用L1L_feat让学生模型的中间特征逼近教师对应层的特征用MSE来算L_attn是注意力层输出分布的约束。我实测下来特征蒸馏比单纯输出蒸馏有效得多。教师最终输出是大量信息的混合体学生容量小硬拟合输出很容易顾此失彼中间特征等于把教师的分析过程逐级喂给学生相当于老师一步步带着做题学生每一步都有人验证走偏不到哪里去。参数上我推荐的起点是α0.1、β0.05。蒸馏权重不是越大越好。α太大学生会拼命去对齐教师的特征图忘记自己本职工作是重建像素结果训练集上的L1 Loss在降但验证集PSNR长时间不动。我后面在问题排查章节会细说这个现象。3.3 第三步重参数化重构把多个分支合并成单分支蒸馏结束后模型结构已经固定了。我再做一次“推理等价变换”这一步是让推理更快的关键。最常见的是ConvBN融合。训练阶段BatchNorm放在卷积后面能稳定梯度分布但推理时多一次归一化运算纯属浪费。卷积是线性操作BN的均值、方差、缩放、偏移都可以反推回卷积核里等价成一个更大的卷积。类似的训练时我可以保留“3x3卷积 1x1卷积 恒等映射”这样的多分支结构让梯度路径更丰富模型更容易收敛推理时把它们合并成单个3x3卷积。y Conv3x3(x) Conv1x1(x) x这个公式在数学上是成立的。1x1卷积的核可以空间补零变成3x3恒等映射等价于一个单位1x1卷积再补零三个相加得到一个新的3x3卷积核。训练时收益是优化更稳推理时收益是算子数量变少。为什么算子数量很重要SwinIR在CPU和低端GPU上跑不快很多时候瓶颈根本不是FLOPs而是算子碎片化。一个Block里十几个小算子来回调度调度开销比计算本身还大。重参数化融合后算子个数变少实测512x512输入的单卡推理耗时大约可以再降5%到8%具体数字和设备相关。3.4 训练配置与完整流程记录整个训练我分了三个阶段每个阶段的职责都不同。第一阶段剪枝后的学生模型随机初始化只用重建损失训练让新结构先把基础能力学到手。第二阶段加入蒸馏损失教师固定α和β用预热方式从0逐渐升到目标值避免训练初期被蒸馏损失带偏。第三阶段去掉蒸馏用小学习率微调重建损失相当于学生独立“上考场”不再依赖教师的中间结果。# Phase 1: baseline student python train.py --arch student --teacher none \ --data_root /data/DF2K --patch_size 128 --batch_size 16 \ --lr 2e-4 --epochs 120 --loss l1 # Phase 2: distillation python train.py --arch student --teacher swinir \ --data_root /data/DF2K --patch_size 128 --batch_size 16 \ --lr 1e-4 --epochs 100 --loss l1 \ --distill 1 --feat_weight 0.1 --attn_weight 0.05 # Phase 3: fine-tune python train.py --arch student --teacher none \ --data_root /data/DF2K --patch_size 128 --batch_size 16 \ --lr 2e-5 --epochs 80 --loss l1训练数据用的DIV2K训练集验证集放在Set5、Set14、BSD100、Urban100上。补丁大小128x128batch size 16优化器AdamW初始学习率2e-4cosine退火前5个epoch做warmup。数据增强是随机翻转和90度旋转。这套配置在RTX3090单卡上大概要跑三到四天如果是8卡可以缩短到一天以内。手头只有单卡的话可以把patch降到96epoch降到80、60、50整体趋势不会变只是最终涨幅会低一些。我自己先用小配置跑通流程确认代码没跑偏再开完整的训练这个习惯帮我避开了好几次浪费一整天的批量训练。4. 实验记录尺寸、精度、速度一起看4.1 评测指标与测试环境评测这件事要单独拿出来说因为很多复现失败都出在评测脚本不统一。我这里统一用PyTorch 1.13 CUDA 11.7RTX 3090单卡输入512x512FP32精度测评时把RGB转成YCbCr后在Y通道上计算PSNR和SSIM完全没有用额外技巧。所有结果都是项目本地跑出来的实测均值不是SwinIR论文抄来的数字硬件不同会有浮动但趋势是可参考的。4.2 总体结果对比下面这组数据是我项目里多次实验取的平均值。参数量和FLOPs是模型本身属性不随机器变化PSNR以本地复现SwinIR为对齐基准推理耗时是同一台机器、同一份脚本测出来的相对值。模型Params (M)FLOPs相对值Set5 PSNRSet14 PSNRBSD100 PSNRUrban100 PSNR推理耗时相对值SwinIR本地复现11.8100%32.9129.0527.8933.15100%本文压缩模型10.486%33.0829.2127.9633.31约90%参数量下降12%FLOPs下降14%四个验证集PSNR全部涨了0.13到0.17dB推理耗时也降了约10%。这个结果说明压缩不一定会牺牲精度方法对路之后小模型反而能通过蒸馏学到更干净的表示。4.3 消融实验每一步的贡献到底有多大光看最终结果不知道每步起了多少作用所以我又跑了一组消融把所有变体控制在同一个评测环境下。配置Set5 PSNR相对SwinIRSwinIR 基线32.91-剪枝后直接训练不加蒸馏32.76-0.15剪枝 输出蒸馏32.90-0.01剪枝 特征蒸馏32.960.05剪枝 特征蒸馏 重参数化重构33.080.17从表格能看得很清楚剪枝本身带来0.15dB的掉点这在预期内输出蒸馏只能把它拉回基线附近特征蒸馏才是真正让模型反超的关键提供了约0.05dB的额外增益。重参数化重构对精度本身影响很小但它通过多分支结构带来了隐式的正则化效果加上后面第三阶段的精调把最终涨幅推到了0.17dB。换句话说剪枝决定模型的下限蒸馏决定上限。只靠剪枝想实现“又小又好”是不现实的一定要把蒸馏的权重调到位。4.4 从可视化看涨点到底长在哪里PSNR是数字落不到眼睛上。我把两个模型在Urban100上的输出做了局部放大重点看三类区域文字边缘、密集栅栏、建筑线条。SwinIR在一些重复纹理区域偶尔会出现振铃细看就是边缘外侧有一圈淡淡的波纹。压缩后的模型在这些区域更干净线条更连续。原因推测有两点一是全局通道剪枝把那些贡献噪声多于有效信息的通道过滤掉了这些通道在原始模型里更像是参数冗余二是特征蒸馏让小模型的中间层学会了教师的主要响应模式丢掉了一些不稳定的旁支响应。文字区域两者的差距最直观小模型的笔画边缘更实没有那种“糊了又试图锐化”的脏感。看结果时建议自己截几组图放大去比只盯PSNR容易忽略局部伪影。5. 压缩落地路上的坑和排查方法5.1 剪枝后PSNR不升反降的常见原因我自己最早一版剪枝实验直接把PSNR干掉了0.3dB以上当时一度以为方法走不通。后来排查下来原因有三个。第一全局剪枝比例拉得太猛。第一次就设了15%以上学生模型骨干通道不够用蒸馏怎么拉都拉不回来。后来把第一轮剪枝比例控制在10%到12%再根据掉点幅度决定要不要往上加。第二残差分支被误剪。剪枝脚本默认对所有卷积统一处理没有给参与残差连接的层加保护导致特征流断裂网络即使重训练也很难恢复。第三重训练周期不够。剪枝后只跑了30个epoch就觉得不行了其实结构突变后模型需要相当长的重新收敛时间至少80个epoch起步。排查口诀可以记一下先看剪枝比例再看残差保护最后看训练长度。如果三者都正常还是掉点0.3dB以上那说明剪枝对象本身选错了比如把短残差里的关键通道剪了需要重新设计剪枝策略。5.2 蒸馏损失权重调不好怎么处理蒸馏损失不是加上就完事我踩过的典型现象是从训练日志看L1 Loss一直在下降但验证PSNR长时间不动甚至开始下降。原因通常是蒸馏损失权重过高学生模型把精力全放在对齐教师特征上丢了自己的重建本职工作。解决方法是先给蒸馏损失加warmup。前20个epoch让α和β从线性增加到目标值给学生一段只学重建的时间再逐步加入教师约束。另一个经验是特征图尺度不一致会导致MSE数值虚高对训练很不稳定。先把教师和学生的特征分别做L2归一化再算MSE训练会平稳非常多。权重的话我的安全区是feat_weight 0.05到0.2attn_weight 0.01到0.05。比如α取0.1、β取0.05是一个不容易出错的起点。不同数据集上最优值会偏移但不会差一个数量级。5.3 模型变小了为什么线上推理时间没变这是一个特别常见的困惑参数和FLOPs都降了部署后延迟却不降。关键原因是推理时间不只由FLOPs决定内存访问次数和算子调度开销经常才是瓶颈。模型参数减少后如果推理框架还在跑几十个碎片化的算子调度器来回切换的时间就占掉了大头。我的建议是按顺序做三步优化。第一步把训练好的网络做重参数化融合能合并的卷积尽量合并减少算子数量。第二步导出ONNX后用TensorRT或OpenVINO这类推理引擎跑让框架自动做算子融合和图优化。第三步如果延迟还差一点再做INT8量化。import torch model load_fused_model() dummy_input torch.randn(1, 3, 512, 512) torch.onnx.export(model, dummy_input, super_res.onnx, opset_version13, input_names[input], output_names[output])量化时校准集一定要使用评测集之外的图片否则量化参数会过拟合到评测集上PTQ之后数字很漂亮换真实数据就露馅。5.4 视频超分场景下要额外注意什么视频超分不是把图像超分模型逐帧跑一遍那么简单模型内部往往多了光流、特征对齐、时序融合这些模块。如果直接用图像超分的剪枝比例去砍很容易砍到时序对齐关键通道轻则出现闪烁重则动态模糊。我建议对光流或对齐模块设置单独的保护规则剪枝目标主要集中在空间特征提取和重建头部分。另外视频超分训练可以考虑时序一致性蒸馏。除了让学生输出接近教师单帧结果还可以约束相邻帧的输出差异接近教师相邻帧的输出差异。这个思路对减少视频闪烁很有用。如果你手头有视频超分模型可以先拿图像超分压缩后的模型做初始化再做时序微调收敛快、效果也稳。6. 从图像超分到视频超分这套思路还能怎么迁移6.1 视频超分模型的“瘦身”空间往往更大视频超分模型的体积普遍比图像超分大不少因为多了时序模块和光流估计。压缩收益也更明显但同时风险更高。我做视频超分压缩时的建议是上来先做模块级profiling搞清楚每个模块的耗时占比和参数量占比再决定往哪剪。一般规律是对齐模块和融合模块耗时高但冗余未必高空间特征提取模块冗余高但耗时未必高。一场模型优化如果只看通信量不看耗时占比很容易把力气花在收益最小的模块上。我自己做过的项目里先profiling再定方案比直接套用固定剪枝比例省了一半以上的试错时间。6.2 想复现这个结果最简路径是什么如果你也想在自己的项目里复现“又小又好”的效果可以按这个清单走。先把SwinIR官方仓库代码跑通生成一份本地基线所有后续对比都基于这份数字不要去抄论文的PSNR。然后写一个通道剪枝脚本第一版只剪10%带上残差保护名单从头训练120个epoch记录掉点幅度。接着加特征蒸馏教师用原始SwinIRα取0.1β取0.05warmup 20个epoch训练100个epoch看能否超过基线。最后做重参数化融合导出ONNX跑推理速度。整个流程走下来如果一切正常你会看到“小12%、PSNR反超0.17dB”这个结论在自己数据上复现出来。第一次跑可能不会一次到位这时前面章节整理的排查清单就是你最好的工具。这次项目结束之后我最大的体会是压缩模型不是把SwinIR塞进一个更小的壳里而是在限制容量的前提下重新教它该关注哪里。剪枝让模型忘掉冗余蒸馏让模型重新记住重点重参数化则把这些记忆高效地固化到推理流程里。这个思路从图像超分迁移到视频超分时依然成立只是需要额外尊重时序模块的脆弱性。目前我还在继续试量化与蒸馏的进一步结合尤其是INT8部署时蒸馏损失该怎么调整等有稳定结论了再单独写一篇。