## 1. Restormer 架构设计与核心模块拆解 Restormer 不是简单地把 ViT 搬到图像恢复任务上套个壳，而是针对高分辨率图像处理的特殊瓶颈做了大量“手术级”优化。我第一次跑通官方代码时就意识到，它真正解决了三个长期困扰图像恢复模型的实际问题：显存爆炸、长程建模低效、高频细节丢失。它的结构不是堆叠Transformer块，而是一套协同工作的子系统。 先说多头转置注意力（MDTA）。你可能熟悉标准的多头自注意力（MHSA），但 MHSA 在图像上直接用有个致命缺陷——计算量随图像尺寸平方增长。一张 256×256 的图，MHSA 的 QK^T 矩阵就有 65536×65536 大小，显存根本扛不住。Restormer 的做法很聪明：它先把输入特征图沿通道维度分组，每组内部做深度卷积（Dconv）提取局部纹理，再对这些局部特征做转置后的注意力操作。所谓“转置”，是指把原本在空间维度上做的注意力，改成在通道维度上聚合信息。你可以把它理解成“让每个通道学会关注哪些其他通道更重要”，而不是让每个像素去算它和所有像素的关系。这样就把 O(HW×HW) 的复杂度降到了 O(C×C)，其中 C 是通道数，通常只有 64 或 128，完全可控。 再来看门控反卷积前馈网络（GDFN）。传统 FFN 就是两层全连接加激活函数，对图像这种强空间结构的数据来说太“扁平”。GDFN 把第一个线性层换成了可学习的反卷积层，能主动扩大感受野；更关键的是引入了门控机制——用一个 sigmoid 激活的分支控制信息流，类似 LSTM 中的遗忘门。我在调试时发现，去掉这个门控，模型在去模糊任务上 PSNR 会掉 0.8dB 以上，尤其在边缘区域出现明显振铃伪影。这说明 GDFN 不只是提升表达能力，更是给模型装了个“开关”，让它能自主决定哪些高频成分该强化、哪些噪声该抑制。 这两个模块组合起来，形成了 Restormer 的“双循环”处理范式：MDTA 负责跨通道的语义关联（比如识别出这是“玻璃反光”还是“雨痕”），GDFN 负责跨空间的结构重建（比如把模糊的窗框重新拉出锐利线条）。它们不像传统 CNN 那样靠堆叠卷积核硬凑感受野，也不像纯 Transformer 那样靠全局注意力硬算关系，而是在通道与空间两个正交维度上分别做高效建模，最后再融合。这种设计让我在复现时少踩了很多坑——比如不用再为显存不够而强行切 patch，也不用担心注意力机制在大图上失效。 ## 2. 官方代码库结构与环境配置实操 swz30/Restormer 这个仓库之所以被广泛采用，不只是因为模型好，更因为它把工程细节打磨得很扎实。我从零开始搭环境时，发现它比很多开源项目更“接地气”：没有花里胡哨的抽象封装，所有数据加载、训练循环、评估逻辑都写在 .py 文件里，变量命名直白，比如 train_img, val_gt 这种一看就懂的名字。但要注意几个容易被忽略的关键点，否则很可能卡在第一步。 首先是依赖版本。官方 README 写着 “PyTorch >= 1.7”，但实际测试下来，1.10.2 是最稳的版本。我试过 1.12 和 2.0，会在 DataLoader 的 num_workers > 0 时出现随机死锁，尤其是 Windows 系统下。CUDA 版本也得匹配，如果你用的是 RTX 4090，别急着装最新版 cudatoolkit，官方预编译的 torch 1.10.2 对应的是 CUDA 11.3，装 11.7 反而会报错找不到 cuBLAS。我的建议是直接用 conda 创建干净环境： ```bash conda create -n restormer python=3.8 conda activate restormer conda install pytorch==1.10.2 torchvision==0.11.3 torchaudio==0.10.2 cudatoolkit=11.3 -c pytorch pip install opencv-python tqdm scikit-image matplotlib ``` 接着是数据目录结构。很多人卡在这里是因为没按官方约定组织文件夹。它不接受任意路径参数，而是硬编码了相对路径。比如做图像去噪任务，你必须把训练集放在 `datasets/SIDD/train/input_crops/` 和 `datasets/SIDD/train/target_crops/` 下，且图片名要严格对应（xxx_input.png 和 xxx_target.png）。我第一次用自己整理的 DND 数据集时，因为把 GT 图片放在了 `gt/` 子目录而非同级目录，训练时 loss 一直不下降，debug 了两天才发现是路径拼错了。后来我写了个小脚本自动校验： ```python import os input_dir = "datasets/SIDD/train/input_crops" target_dir = "datasets/SIDD/train/target_crops" inputs = set([f.split('_')[0] for f in os.listdir(input_dir) if f.endswith('.png')]) targets = set([f.split('_')[0] for f in os.listdir(target_dir) if f.endswith('.png')]) print("Missing in target:", inputs - targets) print("Missing in input:", targets - inputs) ``` 最后是配置文件。官方用 YAML 管理超参，但 `options/train_Restormer.yml` 里的 `batch_size_per_gpu` 不是最终 batch size，而是每个 GPU 上的样本数。如果你有 2 张卡，实际 batch size 是这个值的两倍。我刚开始设成 16，结果 OOM，后来发现单卡显存只够跑 8。还有 `train_patch_size`，这个值直接影响显存占用——设成 128 比 256 能省 60% 显存，但模型性能只掉 0.2dB，属于非常划算的 trade-off。这些细节官方文档没明说，全是我在反复试错中记下来的。 ## 3. 模型加载与推理流程完整实现 复现 Restormer 最常遇到的问题不是模型跑不起来，而是推理结果和预期差很远：颜色偏灰、对比度低、甚至整张图发绿。这往往不是模型问题，而是数据预处理和后处理环节出了偏差。我整理了一套经过多次验证的端到端推理流程，从读图到保存，每个环节都加了注释和容错处理。 第一步是图像读取与归一化。Restormer 训练时用的是 [0,1] 归一化，不是 ImageNet 那套均值方差标准化。所以千万别用 `transforms.Normalize`，直接除以 255 就行。但要注意 OpenCV 默认读 BGR，而 PyTorch 模型期待 RGB： ```python import cv2 import numpy as np import torch def load_image(path): img = cv2.imread(path) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 转 RGB img = img.astype(np.float32) / 255.0 # 归一化到 [0,1] img = torch.from_numpy(img).permute(2, 0, 1) # HWC -> CHW return img.unsqueeze(0) # 增加 batch 维度 input_tensor = load_image('input.png') # shape: [1,3,H,W] ``` 第二步是模型加载。官方 checkpoint 是 `state_dict` 形式，但 key 名带 `module.` 前缀（因为训练时用了 `nn.DataParallel`）。如果你单卡推理，直接 `load_state_dict` 会报错 key 不匹配。解决方案有两个：要么用 `torch.load(..., map_location='cpu')` 后手动 strip 前缀，要么更稳妥地用 `model.load_state_dict(checkpoint['params'], strict=False)` 并忽略不匹配的 key。我推荐后者，因为官方 checkpoint 里有时会多存些训练中间状态。 第三步是推理与后处理。这里最容易出错的是设备迁移和内存管理。Restormer 输入要求是 4D tensor，但很多新手直接把 3D tensor 送进去，结果报 dimension error。另外，GPU 推理后输出 tensor 还在 GPU 上，必须 `.cpu()` 才能转 numpy。后处理还要注意 clip 到 [0,1] 范围，否则可能出现负值或大于 1 的像素： ```python device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = Restormer().to(device) checkpoint = torch.load('pretrained_models/image_denoising.pth', map_location=device) model.load_state_dict(checkpoint['params']) model.eval() with torch.no_grad(): input_tensor = input_tensor.to(device) restored = model(input_tensor) # 输出也是 [1,3,H,W] restored = torch.clamp(restored, 0, 1) # 强制裁剪 restored = restored.cpu().squeeze(0) # 移除 batch 维度并回 CPU restored = restored.permute(1, 2, 0).numpy() # CHW -> HWC restored = (restored * 255).astype(np.uint8) # 转 uint8 cv2.imwrite('output.png', cv2.cvtColor(restored, cv2.COLOR_RGB2BGR)) ``` 我特别强调 `torch.clamp` 这一步。有次我漏了它，输出图里出现大片纯黑区域，查了半天才发现是某些区域预测值小于 0，转 uint8 时溢出成 255 了。这种 bug 很隐蔽，肉眼很难发现，一定要加日志检查输出 tensor 的 min/max 值。 ## 4. 自定义训练与关键参数调优策略 如果你想在自己的数据集上微调 Restormer，或者从头训练一个新任务（比如水下图像增强），官方代码已经提供了完整的训练框架，但有几个参数必须根据你的硬件和数据特性重设，否则大概率训崩。我拿自己在真实监控视频去雾任务上的经验来说，分享一套经过实战检验的调优策略。 首先是学习率调度。官方用的是余弦退火（CosineAnnealingLR），但初始学习率不能照搬 SIDD 数据集的 2e-4。我的监控数据集噪声模式更复杂，初始 lr 设成 1e-4 更稳。更重要的是 warmup 步数——前 500 步先用线性 warmup，避免模型早期震荡太大。我在 `train_Restormer.py` 里加了这段： ```python scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt['train']['n_epoch'] - 500) warmup_scheduler = optim.lr_scheduler.LinearLR(optimizer, start_factor=1e-6, total_iters=500) scheduler = optim.lr_scheduler.SequentialLR(optimizer, schedulers=[warmup_scheduler, scheduler], milestones=[500]) ``` 其次是损失函数选择。官方默认用 L1 Loss，但对去雾这种任务，L1 容易导致结果过平滑。我替换成 Charbonnier Loss（一种平滑的 L1 变体），公式是 √(x² + ε²)，ε 设为 1e-6。它在小梯度时近似线性，在大梯度时接近二次，能更好保留边缘。代码只需改一行： ```python # 原来是 criterion = torch.nn.L1Loss() criterion = lambda x, y: torch.sqrt((x - y) ** 2 + 1e-6).mean() ``` 再就是数据增强。Restormer 训练时默认只做随机翻转和旋转，但对低质量监控数据，我额外加了随机 JPEG 压缩（quality 30~80）和高斯模糊（kernel 3×3，sigma 0.5~1.5），模拟真实压缩伪影。这些增强必须在 `data/derain_dataset.py` 的 `__getitem__` 里实现，不能用 torchvision 的 transforms，因为那些会破坏 tensor 的连续性，导致后续卷积出错。 最后是显存优化技巧。如果显存不够跑 full-size 图像，不要简单缩小 patch size。我试过两种方案：一是用梯度检查点（gradient checkpointing），在 `Restormer` 类的 `forward` 方法里对每个 block 包一层 `torch.utils.checkpoint.checkpoint`，显存能省 40%，速度只慢 15%；二是用混合精度训练（AMP），加三行代码： ```python scaler = torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): loss = criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() ``` 这套组合拳让我在单张 RTX 3090 上，用 256×256 patch 训练去雾模型，batch size 达到 12，比原始配置快 2.3 倍，而且收敛更稳定。关键是要理解每个参数背后的物理意义，而不是盲目调参。