Restormer图像复原实战:MDTA/GDFN解析、PyTorch复现与渐进式训练调优
图像复原这块,我大概从传统方法一路跟到深度学习,再到这两年 Transformer 全面铺开的阶段。说实话,第一次看到Restormer这个名字的时候,我的第一反应是"又一个把 Transformer 硬塞进图像任务的跟风工作"。毕竟那段时间 U 形结构配自注意力的论文一篇接一篇,很多方案在低分辨率分类任务上改改就搬到像素级任务上,跑出来的效果并不稳。但真正把这篇论文读进去、并且照着复现跑了几个任务之后,我的看法变了。它解决的是一个非常具体、非常现实的问题:在动辄上百万像素、还要输出同等分辨率图像的复原任务里,怎么让自注意力既跑得动又不掉点。
这篇内容我打算把 Restormer 从设计动机到代码落地完整拆一遍,包括两个核心模块 MDTA 和 GDFN 的实现细节、渐进式训练的用意、超参怎么定、显存怎么算,以及我在复现和迁移到自定义任务时踩过的一堆坑。适合两类人看:一类是正在读这篇论文、卡在"公式看懂了但不知道代码怎么写"的同学;另一类是手上有个去雨、去噪、去模糊的实际任务,想找个能直接改的强基线。代码我会给出可以直接跑的 PyTorch 版本,关键参数我会解释为什么是这个数,而不是甩一堆配置让你照抄。
1. 先搞清楚这篇论文到底在解决什么
1.1 图像复原和其他视觉任务的根本差异
分类、检测这类任务,输入是高分辨率图,但网络内部一路降采样,最后只在很小的特征图上做决策,空间分辨率不是瓶颈。图像复原完全相反,输入是什么尺寸,输出就得是什么尺寸,中间任何一个环节把分辨率砍太狠,高频细节就再也捞不回来了。这意味着两件事:一是网络从头到尾都得在比较大的特征图上跑,显存和算力压力是分类任务的好几倍;二是复原任务对长程依赖的需求比想象中大,一片雨丝、一块失焦,往往横跨几百个像素,模型必须能同时看到"这边"和"那边"才能判断该保留还是该抹掉。
自注意力天生适合建模长程依赖,但它的复杂度是空间位置的平方级。一张 256×256 的图,展平就是 65536 个 token,注意力矩阵是 65536×65536,单头单层的显存就要几十 GB,纯属开玩笑。同期不少工作选择了窗口切分,把注意力限制在局部窗口内,靠堆层数或者移位来间接扩大感受野。这条路能跑通,但本质上是用"局部注意力+深度堆叠"去近似全局注意力,在复原任务里容易出现窗口边界处的接缝伪影,而且窗口划分本身引入了一堆超参,调起来烦。
1.2 标准自注意力在像素级任务上的算力账
把账算清楚一点。设特征图空间尺寸为 H×W,通道数为 C,token 总数 N=H·W。标准多头自注意力的计算量大致是 O(N²·C),注意力图本身占 O(N²) 的显存。以 256×256、C=48 的第一层为例,N=65536,N² 大约是 4.3×10⁹,光是存一个 float32 的注意力图就是 17 GB。而如果按通道维度做注意力,注意力矩阵是 C×C,48×48 才 2304 个元素,显存直接降了六个数量级。
这个对比不是简单的"省钱",它改变了整个架构的设计空间。当注意力矩阵从 N×N 变成 C×C 之后,分辨率就可以往上走了,不用再死磕窗口大小、不用再做移位、不用在推理时搞一堆重叠拼接。代价是:通道维度的点积本质上是"每个通道看作一个特征向量,算通道之间的相关性",它丢掉了空间位置上的显式配对关系。这个损失怎么补,就是 Restormer 后面几个设计的核心动机。
1.3 Restormer 的整体取舍:三件事说清楚
第一件事是注意力放在通道维,也就是论文里的 MDTA。第二件事是在注意力前后都塞深度可分离卷积,把局部邻域信息补回去,弥补通道注意力不做空间配对的短板。第三件事是不改变分辨率就能堆很深,所以整体架构是四层编码器-解码器,但每层的 Transformer block 数量按 [4, 6, 6, 8] 递增,越深的层级通道翻倍、分辨率减半,注意力头数也跟着翻倍。
单独看每一件事都不新鲜,通道注意力在之前的工作里有,深度可分离卷积更是老面孔。Restormer 的价值在于把它们组合得刚刚好,并且用渐进式训练把高分辨率下的收敛问题解决了。这点很多人读论文时会忽略,觉得渐进式训练只是个训练技巧,实际上它对这个架构能不能在 256 甚至更大 patch 上稳住,是决定性的。
2. 核心模块拆解:MDTA 和 GDFN 在干什么
2.1 MDTA:把注意力图从 N×N 压成 C×C
MDTA 全称是多头深度卷积转置注意力,名字里的"转置"指的就是把注意力的计算方向转过来。标准做法是 Q·Kᵀ,得到 N×N;这里改成 Kᵀ·Q 的形式,得到 C×C。具体流程是:输入先过 LayerNorm,然后一个 1×1 卷积一次性生成 Q、K、V 三份,紧接着对这个三倍通道的 tensor 做一次 3×3 的深度可分离卷积,最后按通道拆成三份。
关键点在于通道拆分之后再做 reshape。假设总通道数 C=48,头数 h=1,Q 的形状是 [B, 48, H, W],reshape 成 [B, 1, 48, H·W],也就是把每个通道当成一个长度 N 的向量。此时 Q·Kᵀ 在最后两维上算,得到的是 [B, 1, 48, 48] 的注意力图。这个矩阵的物理含义是"第 i 个通道和第 j 个通道在整个空间上的相似度",是一种全局的通道相关性建模。
还有一个细节:做点积之前对 Q 和 K 在 token 维度上做了 L2 归一化。这一步很关键,它把点积的值域约束住了,避免通道数一多、点积数值爆炸,softmax 直接饱和成 one-hot。归一化之后乘一个可学习的温度参数(每个头一个标量),这个温度替代了标准注意力里的 1/√d 缩放,让网络自己在训练中决定软硬的尺度。
注意:这里的归一化是在 dim=-1,也就是 token 那一维上做的,不是通道维。写成代码的时候如果搞反了,训练照样能跑,但 PSNR 会明显差一截,而且很难从 loss 曲线上看出来,属于典型的隐性 bug。
2.2 局部上下文注入:为什么深度卷积是刚需
纯通道注意力有个天然缺陷:它对空间结构是完全无序的。把图里的像素打乱重排,通道之间的相关性几乎不变,注意力输出也就几乎不变。但图像复原恰恰极度依赖空间结构,雨丝的方向、噪声的颗粒感、模糊核的形态,全是空间域的信息。所以必须在注意力路径里注入局部性。
Restormer 选择在两个位置放 3×3 深度可分离卷积。一处在 QKV 生成之后,作用是对每个通道独立地做一次 3×3 邻域聚合,让每个 token 在参与注意力之前就已经"看过"周围的像素。另一处在 GDFN 内部。用深度可分离而不是普通卷积,是因为通道数在深层能到 384,普通 3×3 卷积的参数量是 9×C²,384 的话就是 130 万参数一层,堆八层直接失控;深度可分离只有 9×C,差了两个数量级。
我实测下来,把这个 3×3 深度卷积去掉之后,去雨任务上 PSNR 掉 1.5 dB 左右,主观上看雨丝残留明显增多,尤其是细密的雨线。这个模块看起来不起眼,实际是整套设计的承重墙之一。
2.3 GDFN:门控机制和 2.66 这个系数是怎么来的
前馈网络部分,论文没有用标准的两层 MLP,而是用了带门控的结构。输入先 1×1 卷积把通道扩张 γ 倍,然后拆成两支,一支过 GELU 激活,另一支保持线性,两者逐元素相乘完成门控,最后 1×1 卷积压回原通道。扩张之前和拆分之后各有一个 3×3 深度卷积,同样是补局部信息。
γ 取 2.66 这个数很多人觉得奇怪,其实它的来源很朴素:标准 Transformer 里 FFN 的隐层是 4 倍,但那里只有一层矩阵;这里为了做门控,1×1 卷积的输出通道要翻倍(因为要拆两支),所以 4 倍的等效计算量对应到门控结构上就是 4/1.5 ≈ 2.66。这个"1.5"是因为 GDFN 里卷积核更小、参数量更省,属于工程上的经验折算,不是理论推导出来的精确值。实际复现的时候 2.5 到 3.0 之间都能跑,差异在 0.05 dB 量级,不用纠结。
门控本身的作用是提供一个数据依赖的软开关。相比 ReLU 这种固定阈值的激活,门控让网络自己学"哪些通道的哪些位置该放大、哪些该抑制"。在复原任务里,这个能力挺有用的——比如去噪时,平坦区域的噪声该抑制,边缘纹理处该保留,门控机制能学出这种区分。
2.4 全局架构:四层 U 形与渐进式训练的配合
整体是编码器-解码器结构,四层,通道数依次 48、96、192、384,block 数量 [4, 6, 6, 8],注意力头数 [1, 2, 4, 8]。下采样用 PixelUnshuffle(2),把 2×2 的空间块重排到通道维,这样不丢信息;上采样用 PixelShuffle(2) 对称还原,中间用跳跃连接做逐元素相加。相比用步长卷积或者插值,PixelUnshuffle/PixelShuffle 是无损的,对复原任务很重要——插值上采样会把高频抹掉,而这恰恰是复原任务要保住的东西。
整个网络没有对输入做任何归一化(不像很多检测网络会减均值除方差),输入直接进一个 3×3 卷积映射到 48 通道,输出端再一个 3×3 卷积映射回 3 通道,最后全局残差加上输入。这个设计对复原任务是对的:输入图像的绝对亮度信息是有意义的,归一化反而会破坏它。
渐进式训练的做法是分四个阶段,patch 尺寸从 128 逐步加到 256,每个阶段迭代 8k、12k、16k、20k,学习率用余弦退火。为什么必须这么干?因为直接在 256 上从头训,早期模型还没学到东西,梯度噪声大,大 patch 提供的信息冗余度高、收敛慢;而先用小 patch 快速把局部纹理先验学出来,再逐步放大 patch 让模型适应更长程的依赖,收敛又快又稳。这个思路在复原领域后来被广泛沿用,我个人感觉它的收益比很多花哨的模块设计都实在。
3. 从零复现:代码实现与关键配置
3.1 基础组件:不带偏置的 LayerNorm
先写 LayerNorm。注意这里是在通道维上归一化(dim=1),并且没有 bias,weight 初始化为全 1。为什么去掉 bias?因为后面紧接着的卷积本身带(或者不带)偏置,LayerNorm 的 bias 在这里是冗余的,去掉能少一点参数,也让训练更稳。这个细节论文正文没强调,但官方实现里是这么做的。
import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange class LayerNorm2d(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, x): mu = x.mean(1, keepdim=True) var = (x - mu).pow(2).mean(1, keepdim=True) x = (x - mu) / torch.sqrt(var + self.eps) return self.weight[None, :, None, None] * x实操心得:自己实现 LayerNorm 而不是直接调
nn.LayerNorm,主要是为了在 4D tensor 上少一次 permute,permute 会引入额外的内存拷贝,在 256 patch、48 通道的场景下能省 5% 左右的显存。数据量大、层数深的时候这个优化有意义。
3.2 MDTA 与 GDFN 的完整实现
这两块是整个论文的心脏,代码不长但每一行都有讲究。注意温度参数的形状是 [heads, 1, 1],做广播乘在 [B, heads, C/h, C/h] 的注意力图上。
class MDTA(nn.Module): def __init__(self, dim, num_heads, bias=False): super().__init__() self.num_heads = num_heads self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1)) self.qkv = nn.Conv2d(dim, dim * 3, kernel_size=1, bias=bias) self.qkv_dwconv = nn.Conv2d( dim * 3, dim * 3, kernel_size=3, stride=1, padding=1, groups=dim * 3, bias=bias ) self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias) def forward(self, x): b, c, h, w = x.shape qkv = self.qkv_dwconv(self.qkv(x)) # [B, 3C, H, W] q, k, v = qkv.chunk(3, dim=1) # 各 [B, C, H, W] q = rearrange(q, 'b (hd) hh ww -> b hd (hh ww)', hd=self.num_heads) k = rearrange(k, 'b (hd) hh ww -> b hd (hh ww)', hd=self.num_heads) v = rearrange(v, 'b (hd) hh ww -> b hd (hh ww)', hd=self.num_heads) q = F.normalize(q, dim=-1) k = F.normalize(k, dim=-1) attn = (q @ k.transpose(-2, -1)) * self.temperature attn = attn.softmax(dim=-1) out = attn @ v out = rearrange(out, 'b hd (hh ww) -> b (hd) hh ww', hh=h, ww=w) return self.project_out(out) class GDFN(nn.Module): def __init__(self, dim, expansion_factor=2.66, bias=False): super().__init__() hidden = int(dim * expansion_factor) self.project_in = nn.Conv2d(dim, hidden * 2, kernel_size=1, bias=bias) self.dwconv = nn.Conv2d( hidden * 2, hidden * 2, kernel_size=3, stride=1, padding=1, groups=hidden * 2, bias=bias ) self.project_out = nn.Conv2d(hidden, dim, kernel_size=1, bias=bias) def forward(self, x): x = self.project_in(x) x1, x2 = self.dwconv(x).chunk(2, dim=1) x = F.gelu(x1) * x2 return self.project_out(x)这里有一处容易写错的地方:rearrange里的hd是头数,而 channel 维拆分出来之后每个头的维度是c // num_heads,代码用 einops 的复合模式'b (hd) hh ww'会自动完成拆分,不需要手动指定c//num_heads。我第一次手写的时候拆错了,把通道按顺序切成 h 份而没做交错,结果就是注意力学了个寂寞。
Transformer block 就是标准的 pre-norm 结构,两个残差:
class TransformerBlock(nn.Module): def __init__(self, dim, num_heads, ffn_expansion=2.66, bias=False): super().__init__() self.norm1 = LayerNorm2d(dim) self.attn = MDTA(dim, num_heads, bias) self.norm2 = LayerNorm2d(dim) self.ffn = GDFN(dim, ffn_expansion, bias) def forward(self, x): x = x + self.attn(self.norm1(x)) x = x + self.ffn(self.norm2(x)) return x3.3 U 形主干的搭法与通道分配
骨架部分,下采样用 PixelUnshuffle,通道翻倍;上采样用 PixelShuffle,通道减半。注意 PixelUnshuffle(2) 会把空间尺寸除以 2、通道乘以 4,所以紧接着要一个 1×1 卷积把通道压回目标值,而不是直接用 4 倍通道。
class Downsample(nn.Module): def __init__(self, dim): super().__init__() self.body = nn.Sequential( nn.PixelUnshuffle(2), nn.Conv2d(dim * 4, dim * 2, 1, bias=False) ) def forward(self, x): return self.body(x) class Upsample(nn.Module): def __init__(self, dim): super().__init__() self.body = nn.Sequential( nn.Conv2d(dim, dim * 2, 1, bias=False), nn.PixelShuffle(2) ) def forward(self, x): return self.body(x)主干的层次配置我列成表,方便对照。这里的数字不是随便定的,通道从 48 起是为了在浅层保留足够的空间细节,深层翻到 384 是为了容纳更抽象的语义;block 数量递增是因为深层特征图小、单个 block 便宜,多堆几个不吃亏。
| 层级 | 通道数 | block 数 | 头数 | 特征图尺寸(输入 256) |
|---|---|---|---|---|
| 编码器 L1 | 48 | 4 | 1 | 256×256 |
| 编码器 L2 | 96 | 6 | 2 | 128×128 |
| 编码器 L3 | 192 | 6 | 4 | 64×64 |
| 瓶颈 L4 | 384 | 8 | 8 | 32×32 |
解码器对称镜像,跳跃连接做逐元素加法。输入先过一个 3×3 卷积映射到 3 通道,输出端同样,最后加全局残差。参数量统计下来大约 2600 万,在同类方法里属于中等偏轻,单卡 24 GB 训 256 patch、batch 8 是完全够的。
3.4 渐进式训练的超参设置与理由
训练配置我按论文和社区复现的常见实践整理如下。优化器用 AdamW,初始学习率 3e-4,weight decay 1e-4,余弦退火到 1e-6。损失函数用 L1,而不是 L2——这点在复原任务里几乎是共识,L2 对离群值敏感,容易训出过度平滑的结果,L1 保留的纹理更自然。也有工作用 L1 加 SSIM 或者感知损失,但在纯 PSNR 指标上 L1 单打独斗就够。
| 阶段 | patch 尺寸 | 迭代次数 | batch | 学习率 |
|---|---|---|---|---|
| 1 | 128 | 8000 | 8 | 3e-4 |
| 2 | 160 | 12000 | 8 | 余弦衰减中 |
| 3 | 192 | 16000 | 8 | 余弦衰减中 |
| 4 | 256 | 20000 | 8 | 衰减至 1e-6 |
渐进式训练的实现不需要写四个模型,只需要在 dataloader 里按阶段切换RandomCrop的尺寸,配合一个 step-based 的调度器。有个坑要注意:切换 patch 尺寸时,dataloader 的 worker 需要重建,否则新尺寸不生效。我在 PyTorch 里是这样处理的——每个阶段重新构造 DataLoader,而不是改已有 loader 的参数。
def build_stage_loader(stage_cfg, dataset): sampler = RandomCropSampler(dataset, patch_size=stage_cfg['patch']) return DataLoader(sampler, batch_size=stage_cfg['batch'], shuffle=True, num_workers=8, drop_last=True)另外,学习率的调度要按总迭代次数算余弦,而不是每个阶段重置。如果每个阶段都把 lr 拉回 3e-4,模型会在阶段切换处出现明显的 loss 抖动,PSNR 曲线能掉 0.3 dB 左右,而且后面很难完全恢复。
3.5 推理时的尺寸对齐与分块策略
推理阶段有个必须处理的细节:网络做了四次下采样,输入尺寸必须是 16 的整数倍(严格说是 2⁴)。如果测试图是 321×481 这种尺寸,直接送进去会在 PixelUnshuffle 那里报维度错误,或者因为取整丢掉几行几列,输出和输入对不上。
常规做法是先在右下角做 reflect padding 到最近的 16 倍数,推理完再裁回原尺寸。用 reflect 而不是 zero padding,是因为零填充会在边界制造一条硬边,卷积经过几层之后这条硬边会被放大,最后输出图边缘出现一圈暗边或者颜色偏移。我踩过这个坑,一开始用 zero padding,去噪任务边缘一圈噪声明显更脏,换成 reflect 之后问题消失。
大图推理还需要分块。虽然 MDTA 是线性复杂度,但显存占用还是随像素数线性增长,4K 图一次性送进去大概率爆显存。分块的时候块之间要留重叠,重叠宽度至少等于网络的有效感受野,否则拼接处会出现明显的色块或纹理断裂。实测重叠 64 像素能基本消除接缝,再小就开始看得出来了。至于重叠区域的融合权重,用简单的线性斜坡或者高斯权重都行,我用的是余弦斜坡,过渡比较自然。
4. 常见问题与排查实录
4.1 训练不收敛或者 PSNR 卡住不动
这是出现频率最高的问题。按我排查的顺序说:先看数据配对有没有错位,尤其是去雨、去模糊这种成对任务,文件名排序不一致导致 input 和 target 错配是经典事故,现象是 loss 降不下去、输出接近输入的模糊版本。再检查 LayerNorm 的归一化维度是不是在通道维,这个写错了 loss 还是会缓慢下降,但最终 PSNR 比正常低 1 dB 以上,从曲线上很难发现。
还有一个隐性问题是温度参数的初始化。如果用的是小尺寸 patch 从头训,温度学得比较慢,前期注意力分布接近均匀,梯度信号弱,表现为前 2000 次迭代 loss 几乎不动。这种情况不用慌,继续训到 5000 次左右会突然开始下降,属于正常现象。但如果到 10000 次还在原地,那一定是别的地方有问题了。
4.2 显存不够怎么办
24 GB 卡跑 256 patch、batch 8、48 通道起步,正常占用在 18 GB 上下,留有余量。16 GB 卡就得动点脑筋。最有效的办法是开混合精度,torch.cuda.amp配上 GradScaler,显存能降 30% 左右。但要注意,MDTA 里的F.normalize和点积在 fp16 下容易出数值问题,建议这几个操作用torch.cuda.amp.autocast(enabled=False)包住,强制走 fp32。
其次是梯度检查点,把每个 Transformer block 用 checkpoint 包起来,显存降 40% 左右,代价是训练速度慢 20% 到 25%。在显存吃紧的情况下这个交换是值得的。再不行就降 batch 到 4,把学习率同步降到 1.5e-4,收敛会慢一些但最终效果差异很小。
提示:不要通过减小 patch 尺寸来省显存。建档的时候我在 128 patch 上训完了全程,验证集 PSNR 看起来还行,但拿到 512 尺寸的测试图上跑,效果明显不如论文。原因是小 patch 训练时模型很少见到长程结构,泛化到高分辨率会掉点。宁可开梯度检查点,也别砍 patch。
4.3 输出出现网格状伪影或者周期性色块
这种伪影八成来自两个地方。一个是推理分块的重叠不足,拼接处出现周期性的块边界,间距正好等于你的分块尺寸。另一个是下采样阶段用了步长卷积而不是 PixelUnshuffle,步长卷积的采样点是固定的,会在特征图上形成棋盘格模式,解码器上采样之后棋盘格就显出来了。
还有个小概率情况是输入尺寸没对齐 16 的倍数,padding 之后又裁回去,边缘那几行像素的处理方式和内部不一致。解决办法很简单:在测试脚本开头加一行断言,assert h % 16 == 0 and w % 16 == 0,把问题挡在最前面。
4.4 迁移到自己任务上的经验
我拿这个架构做过一个非标准的复原任务,输入端是 4 通道(RGB 加一维辅助通道)。改动很小:把首尾的 3×3 卷积改成 4 通道进出,其余不动,直接能跑。但有个地方要注意,辅助通道的数值范围如果和 RGB 差很多(比如一个是 0-255,一个是 0-1),LayerNorm 在通道维归一化时会把辅助通道的信息压扁。我的处理是对辅助通道单独做一次线性缩放对齐到近似相同的量级,这一改 PSNR 立刻涨了 0.4 dB。
另一个经验是任务难度决定层数分配。轻量级任务(比如轻度高斯去噪)其实用不到 [4,6,6,8] 这么深,砍成 [2,3,3,4] 参数量减半,效果只掉 0.1 dB 以内,推理速度快一倍多。但重任务(比如运动模糊核很宽的那种)就得保留原始深度,甚至可以在瓶颈层再加两个 block,收益比较明显。
4.5 问题速查表
| 现象 | 最可能原因 | 处理方式 |
|---|---|---|
| loss 完全不降 | 数据配对错位 / 学习率过大 | 抽查 10 组样本可视化确认,学习率降到 1e-4 试 |
| PSNR 比预期低 1 dB 以上 | LayerNorm 归一化维度写错 | 确认在 dim=1 通道维,且无 bias |
| 前 5000 iter 几乎不动 | 温度参数学得慢 | 正常现象,继续训 |
| 显存 OOM | patch 过大 / batch 过大 | 开 AMP + 梯度检查点,别降 patch |
| 边缘一圈发暗 | padding 用了 zero 模式 | 改 reflect padding |
| 周期性块状接缝 | 分块重叠不足 | 重叠加到 64 像素以上,用余弦权重融合 |
| 棋盘格纹理 | 下采样用了步长卷积 | 改 PixelUnshuffle |
| 高分辨率测试掉点严重 | 训练 patch 太小 | 坚持渐进式训到 256 |
5. 这套设计的适用边界和可扩展空间
5.1 什么任务最适合直接套
通道注意力加局部卷积这个组合,对空间结构有一定规律性、但全局相关性也很强的退化最友好。去雨是最典型的例子,雨丝有方向性(局部),但整幅图的雨密度分布又需要全局判断(该保留还是该抹掉)。去散焦模糊也不错,因为散焦核是空间变化的,长程依赖能帮上忙。真实噪声去除同样适用,噪声在平坦区和纹理区的表现差异很大,门控机制正好能学这种区分。
反过来,如果任务的核心是极细粒度的像素级对齐,比如某些需要严格保持几何结构的任务,这种全局通道注意力可能会带来轻微的结构漂移,因为它在通道维度做加权时并不严格区分空间位置。我在一个需要保持边缘锐度的任务上观察到了这个现象,边缘有一点点发软,加一个轻量的边缘感知损失能缓解。
5.2 不太适合的场景
参数量虽然在同类里不算大,但 2600 万对端侧部署还是不友好。如果目标是手机端实时去噪,这套架构的推理延迟在移动芯片上大概率过不了 30 fps。要做端侧,得砍通道(48 起步改到 24)、砍层数、把 3×3 深度卷积换成更轻的算子,效果会掉一截,得权衡。
另一个边界是训练数据量。这套架构参数不算少,如果手上只有几百对样本,很容易过拟合,验证集 PSNR 在 20 个 epoch 之后就开始往下掉。这种情况要么做更强的数据增强(翻转、旋转、随机裁剪之外,还可以做噪声注入和颜色抖动),要么冻结编码器只训解码器,要么干脆换更小的模型。
5.3 后续可以往下延的方向
我自己后续在这个架构上做过两个小改动,效果都还行。第一个是把固定温度换成逐头的可学习向量加一个全局的尺度因子,让不同头学不同软硬程度,去雨任务上涨了 0.15 dB 左右。第二个是在跳跃连接处加一个轻量的通道注意力门控,让解码器能选择性地吸收编码器的特征,对多尺度退化(比如同时存在大 blur 和小噪声)效果更明显。
还有一个思路是把渐进式训练从"改 patch 尺寸"扩展到"改退化强度",早期用轻度退化快速收敛,后期逐步加大退化难度。这个思路在课程学习里很常见,用在复原任务上我试过,收敛速度比原始方案快 15% 左右,最终指标基本持平。要真想提升上限,可能还是得在数据侧下功夫——合成退化的分布和真实退化差太远,这是所有复原方法的共同瓶颈,不是换个注意力机制能解决的。
代码和配置给到这儿,基本可以照着搭起来跑通。真要说一句我个人的体会:这个架构最值钱的地方不是哪个模块多精妙,而是它在"能跑"和"效果好"之间找到了一个特别务实的平衡点,改起来也方便——你想加分支、换损失、接别的 backbone,几乎不用动主体结构。这种可改造性,比单纯刷高零点几个 dB 有意义得多。