简介本资源面向深度学习图像去雾方向的研究者与工程实践者提供一套基于Vision Transformer架构的完整去雾算法实现覆盖NH-HAZE、NTIRE2019、I-HAZE、O-HAZE四种真实雾霾数据集的训练与测试流程。压缩包共342个文件约156.42MB以204个Python脚本为核心配合16个yaml配置、9个ipynb实验笔记、12个csv指标记录及39张png与10个gif可视化结果另含md说明、sh运行脚本与少量文档便于复现与二次开发。资源围绕自注意力机制建模长距离依赖包含数据预处理、模型结构调参、损失函数设计及PSNR、SSIM等量化评估环节并附实验设置说明与结果分析材料。已有125人学习下载适合希望系统掌握Transformer去雾方案、对比传统方法差异并快速搭建实验环境的中高级读者。1. 真实雾霾去雾为什么让 ViT 架构重新洗牌如果你拿合成雾霾数据集比如 RESIDE 里的 ITS/OTS训过一个去雾模型指标刷到 35dB 以上兴冲冲拿真实雾霾照片一测结果大概率是翻车的——颜色发灰、细节糊成一团、远处物体直接消失。这不是模型训崩了而是合成雾霾和真实雾霾之间存在巨大的域 gap。NTIRE2019 去雾挑战赛和后来的 NH-HAZE、I-HAZE、O-HAZE 这几个数据集就是专门用来暴露这个问题的。而 Vision Transformer 架构之所以在这几个真实雾霾数据集上被反复验证核心原因是它的全局注意力机制能建模非局部的雾霾分布不像 CNN 那样受限于局部感受野。这篇笔记面向的是已经跑通过基础去雾模型、想切到真实雾霾数据集上做训练和测试的从业者从数据准备、模型搭建、训练参数到踩坑排查一步步把这条链路走通。2. 四个真实雾霾数据集NH-HAZE、NTIRE2019、I-HAZE、O-HAZE 怎么选怎么用2.1 四个数据集的本质差异与选型逻辑先把这四个数据集的关系理清楚不然后面训练策略会走偏。I-HAZE 和 O-HAZE 是同一套采集体系下的产物I-HAZE 是室内场景用专业雾霾发生器在可控环境下拍摄雾霾浓度均匀、光照稳定O-HAZE 是室外场景同样用雾霾机生成但受自然光影响雾霾分布不均匀。这两个数据集的配对图像有雾/无雾是在同一位置、同一光照下先后拍摄的所以 ground truth 非常干净适合做 baseline 验证。NH-HAZE 是非均匀雾霾数据集同样是室内可控环境但雾霾发生器做了改造能产生浓度空间变化的雾霾。这个数据集更接近真实场景中雾霾分布不均的情况是检验模型泛化能力的关键。NTIRE2019 是挑战赛数据集包含室内和室外两部分图像分辨率高、场景多样但部分图像的 ground truth 对齐存在微小误差训练时需要注意。选型建议很直接如果只是验证算法可行性从 I-HAZE 或 O-HAZE 入手数据干净、训练稳定如果要发论文或做产品级验证NH-HAZE 是必选项因为它最能暴露模型在非均匀雾霾下的短板NTIRE2019 适合做最终的大规模评测。数据集场景雾霾类型图像数量训练/测试分辨率主要用途I-HAZE室内均匀30/5约 2000×1500Baseline 验证O-HAZE室外均匀35/5约 2000×1500室外泛化验证NH-HAZE室内非均匀50/5约 2000×1500非均匀雾霾核心评测NTIRE2019室内室外混合55/10不等大规模综合评测注意这四个数据集的图像数量都很少I-HAZE 只有 30 对训练图直接训 ViT 几乎必然过拟合。后面会讲怎么用数据增强和迁移学习来缓解。2.2 数据加载与预处理的标准写法真实雾霾数据集的目录结构通常不统一我一般会先写一个统一的 Dataset 类把四个数据集的加载逻辑收敛到同一个接口。下面是一个可直接复用的 PyTorch Dataset 实现import os import random from PIL import Image from torch.utils.data import Dataset import torchvision.transforms as T class DehazeDataset(Dataset): def __init__(self, root_dir, modetrain, patch_size256, augmentTrue): root_dir: 数据集根目录内部结构为 root_dir/hazy/ 和 root_dir/clear/ mode: train 或 test patch_size: 训练时随机裁剪的 patch 大小 augment: 是否启用数据增强 self.hazy_dir os.path.join(root_dir, hazy) self.clear_dir os.path.join(root_dir, clear) self.mode mode self.patch_size patch_size self.augment augment # 按文件名排序确保 hazy 和 clear 一一对应 self.filenames sorted(os.listdir(self.hazy_dir)) # 基础变换转 Tensor归一化到 [0,1] self.to_tensor T.ToTensor() def __len__(self): return len(self.filenames) def __getitem__(self, idx): fname self.filenames[idx] hazy Image.open(os.path.join(self.hazy_dir, fname)).convert(RGB) clear Image.open(os.path.join(self.clear_dir, fname)).convert(RGB) if self.mode train: # 随机裁剪同一位置 w, h hazy.size if w self.patch_size or h self.patch_size: # 图像太小则先 resize hazy hazy.resize((max(w, self.patch_size), max(h, self.patch_size))) clear clear.resize((max(w, self.patch_size), max(h, self.patch_size))) w, h hazy.size x random.randint(0, w - self.patch_size) y random.randint(0, h - self.patch_size) hazy hazy.crop((x, y, x self.patch_size, y self.patch_size)) clear clear.crop((x, y, x self.patch_size, y self.patch_size)) if self.augment: # 随机水平翻转 if random.random() 0.5: hazy hazy.transpose(Image.FLIP_LEFT_RIGHT) clear clear.transpose(Image.FLIP_LEFT_RIGHT) # 随机垂直翻转 if random.random() 0.5: hazy hazy.transpose(Image.FLIP_TOP_BOTTOM) clear clear.transpose(Image.FLIP_TOP_BOTTOM) # 随机旋转 90 度的倍数 if random.random() 0.5: angle random.choice([90, 180, 270]) hazy hazy.rotate(angle) clear clear.rotate(angle) return self.to_tensor(hazy), self.to_tensor(clear)这段代码的关键点有三个第一hazy 和 clear 必须用相同的随机种子做裁剪和翻转否则配对关系就断了第二patch_size 设为 256 是 ViT 类模型的常见选择因为 ViT 的 patch embedding 通常要求输入尺寸能被 patch 大小整除第三数据增强只用了几何变换没有用颜色抖动因为颜色抖动会改变雾霾的浓度分布反而引入噪声。参数说明patch_size 在训练时建议 256测试时用原图或 resize 到 512augment 在 I-HAZE 这种小数据集上必须开在 NTIRE2019 上可以酌情关掉一部分。3. Vision Transformer 去雾模型怎么搭从 patch embedding 到重建头3.1 ViT 用于去雾的架构选型理由标准 ViT 是为分类任务设计的直接拿来做去雾有两个问题一是输出是类别概率不是图像二是 patch embedding 会丢失空间细节。所以去雾用的 ViT 需要做三处改造把分类头换成图像重建头、在 encoder 和 decoder 之间加 skip connection、用多尺度 patch embedding 来保留高频信息。我一般会采用类似 Uformer 或 Restormer 的思路encoder 用 ViT block 做全局建模decoder 用轻量卷积做上采样和细节恢复中间用 skip connection 把 encoder 的多尺度特征传到 decoder。这样既有 Transformer 的全局建模能力又有 CNN 的局部细节恢复能力。具体到 ViT block 的设计核心是 Multi-Head Self-Attention 的窗口划分策略。全局注意力计算量太大256×256 的图有 65536 个像素直接做全局注意力显存扛不住。常见做法是分成 8×8 或 16×16 的窗口在窗口内做自注意力再周期性做窗口移位来跨窗口交互。这就是 Swin Transformer 的核心思想也是目前去雾任务里最稳的 ViT 变体。3.2 一个可跑通的 ViT 去雾模型实现下面是一个简化但可运行的 ViT 去雾模型基于 Swin Transformer 的窗口注意力机制代码结构清晰方便你按自己的需求改import torch import torch.nn as nn import torch.nn.functional as F class WindowAttention(nn.Module): 窗口内多头自注意力 def __init__(self, dim, window_size, num_heads): super().__init__() self.dim dim self.window_size window_size self.num_heads num_heads head_dim dim // num_heads self.scale head_dim ** -0.5 # 相对位置偏置 self.relative_position_bias_table nn.Parameter( torch.zeros((2 * window_size - 1) * (2 * window_size - 1), num_heads) ) nn.init.trunc_normal_(self.relative_position_bias_table, std0.02) self.qkv nn.Linear(dim, dim * 3) self.proj nn.Linear(dim, dim) def forward(self, x): B_, N, C x.shape qkv self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads) qkv qkv.permute(2, 0, 3, 1, 4) q, k, v qkv[0], qkv[1], qkv[2] q q * self.scale attn (q k.transpose(-2, -1)) # 这里简化处理实际使用时需要加上相对位置偏置的索引计算 attn attn.softmax(dim-1) x (attn v).transpose(1, 2).reshape(B_, N, C) x self.proj(x) return x class SwinBlock(nn.Module): 一个 Swin Transformer block窗口注意力 FFN def __init__(self, dim, window_size8, num_heads4, shift_size0): super().__init__() self.dim dim self.window_size window_size self.shift_size shift_size self.norm1 nn.LayerNorm(dim) self.attn WindowAttention(dim, window_size, num_heads) self.norm2 nn.LayerNorm(dim) self.ffn nn.Sequential( nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim) ) def forward(self, x): # x: (B, H, W, C) B, H, W, C x.shape shortcut x x self.norm1(x) # 窗口划分 x x.view(B, H // self.window_size, self.window_size, W // self.window_size, self.window_size, C) x x.permute(0, 1, 3, 2, 4, 5).contiguous() x x.view(-1, self.window_size * self.window_size, C) # 窗口注意力 x self.attn(x) # 窗口还原 x x.view(B, H // self.window_size, W // self.window_size, self.window_size, self.window_size, C) x x.permute(0, 1, 3, 2, 4, 5).contiguous() x x.view(B, H, W, C) x shortcut x x x self.ffn(self.norm2(x)) return x class ViTDehaze(nn.Module): 基于 Swin 窗口注意力的去雾模型 def __init__(self, in_chans3, embed_dim64, depths[2, 2, 4, 2], num_heads[2, 4, 8, 16]): super().__init__() self.embed_dim embed_dim # 浅层特征提取 self.conv_first nn.Conv2d(in_chans, embed_dim, 3, 1, 1) # Encoder: 4 个 stage每个 stage 后下采样 self.encoders nn.ModuleList() self.downsamples nn.ModuleList() for i in range(len(depths)): stage nn.Sequential(*[ SwinBlock(embed_dim * (2 ** i), window_size8, num_headsnum_heads[i]) for _ in range(depths[i]) ]) self.encoders.append(stage) if i len(depths) - 1: self.downsamples.append( nn.Conv2d(embed_dim * (2 ** i), embed_dim * (2 ** (i 1)), 2, 2) ) # Decoder: 上采样 卷积 self.decoders nn.ModuleList() self.upsamples nn.ModuleList() for i in range(len(depths) - 1, 0, -1): self.upsamples.append( nn.ConvTranspose2d(embed_dim * (2 ** i), embed_dim * (2 ** (i - 1)), 2, 2) ) self.decoders.append( nn.Sequential( nn.Conv2d(embed_dim * (2 ** (i - 1)) * 2, embed_dim * (2 ** (i - 1)), 3, 1, 1), nn.GELU(), nn.Conv2d(embed_dim * (2 ** (i - 1)), embed_dim * (2 ** (i - 1)), 3, 1, 1) ) ) # 重建头 self.reconstruct nn.Sequential( nn.Conv2d(embed_dim, embed_dim, 3, 1, 1), nn.GELU(), nn.Conv2d(embed_dim, in_chans, 3, 1, 1) ) def forward(self, x): # 浅层特征 feat self.conv_first(x) # Encoder enc_feats [] for i, encoder in enumerate(self.encoders): # 转成 (B, H, W, C) 给 SwinBlock B, C, H, W feat.shape feat feat.permute(0, 2, 3, 1).contiguous() feat encoder(feat) feat feat.permute(0, 3, 1, 2).contiguous() enc_feats.append(feat) if i len(self.downsamples): feat self.downsamples[i](feat) # Decoder for i, (upsample, decoder) in enumerate(zip(self.upsamples, self.decoders)): feat upsample(feat) # Skip connection skip enc_feats[-(i 2)] feat torch.cat([feat, skip], dim1) feat decoder(feat) # 重建 out self.reconstruct(feat) # 残差学习预测雾霾残差而不是直接预测清晰图 return torch.clamp(x - out, 0, 1)这段代码的核心设计思路是残差学习模型不直接输出清晰图像而是输出雾霾层然后用输入减去雾霾层得到清晰图。这样做的好处是训练更稳定因为雾霾层的数值范围比清晰图小梯度更容易控制。参数说明embed_dim 控制基础通道数64 是显存和效果的平衡点depths 控制每个 stage 的 block 数量[2,2,4,2] 是 Swin-T 的配置显存不够可以降到 [2,2,2,2]num_heads 随 stage 递增因为通道数在翻倍。注意这个实现省略了相对位置偏置的完整索引计算和 shift window 的 roll 操作实际训练时需要补上否则窗口之间没有信息交互效果会打折扣。4. 训练参数怎么设学习率、损失函数、batch size 的血泪经验4.1 损失函数组合与权重分配去雾任务的损失函数不能只用 L1 或 L2因为像素级损失会导致结果过于平滑丢失纹理细节。我一般用 L1 Perceptual Loss SSIM Loss 的组合import torch import torch.nn as nn import torchvision.models as models class DehazeLoss(nn.Module): def __init__(self, w_l11.0, w_perceptual0.1, w_ssim0.5): super().__init__() self.w_l1 w_l1 self.w_perceptual w_perceptual self.w_ssim w_ssim # 用 VGG16 的浅层特征做感知损失 vgg models.vgg16(pretrainedTrue).features[:16].eval() for param in vgg.parameters(): param.requires_grad False self.vgg vgg self.l1 nn.L1Loss() def forward(self, pred, target): # L1 损失 loss_l1 self.l1(pred, target) # 感知损失 pred_feat self.vgg(pred) target_feat self.vgg(target) loss_perceptual self.l1(pred_feat, target_feat) # SSIM 损失简化版实际建议用 pytorch_msssim 库 loss_ssim 1 - self.ssim(pred, target) return (self.w_l1 * loss_l1 self.w_perceptual * loss_perceptual self.w_ssim * loss_ssim) def ssim(self, x, y): # 简化 SSIM 计算实际建议用现成库 C1 0.01 ** 2 C2 0.03 ** 2 mu_x F.avg_pool2d(x, 3, 1, 1) mu_y F.avg_pool2d(y, 3, 1, 1) sigma_x F.avg_pool2d(x ** 2, 3, 1, 1) - mu_x ** 2 sigma_y F.avg_pool2d(y ** 2, 3, 1, 1) - mu_y ** 2 sigma_xy F.avg_pool2d(x * y, 3, 1, 1) - mu_x * mu_y ssim_map ((2 * mu_x * mu_y C1) * (2 * sigma_xy C2)) / \ ((mu_x ** 2 mu_y ** 2 C1) * (sigma_x sigma_y C2)) return ssim_map.mean()权重分配的经验值w_l11.0 是基础w_perceptual0.1 不能太大否则颜色会偏w_ssim0.5 用来保结构。在 NH-HAZE 上训练时感知损失的权重可以降到 0.05因为非均匀雾霾下 VGG 特征本身就不稳定。4.2 学习率调度与优化器配置ViT 类模型对学习率非常敏感用 AdamW 比 Adam 稳初始学习率设 2e-4 到 5e-4配合 cosine annealing 衰减。batch size 在单卡 24G 显存下patch_size256 时最多开到 8再大就 OOM。from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR # 优化器 optimizer AdamW(model.parameters(), lr3e-4, weight_decay1e-4) # 学习率调度先 warmup 再 cosine 衰减 def lr_lambda(epoch): if epoch 5: return epoch / 5 # warmup else: return 0.5 * (1 math.cos(math.pi * (epoch - 5) / (total_epochs - 5))) scheduler torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)warmup 阶段很重要ViT 的注意力层在训练初期梯度很大直接上大学习率会震荡。5 个 epoch 的 warmup 在四个数据集上都比较稳。4.3 训练循环与验证策略for epoch in range(total_epochs): model.train() for hazy, clear in train_loader: hazy, clear hazy.cuda(), clear.cuda() pred model(hazy) loss criterion(pred, clear) optimizer.zero_grad() loss.backward() # 梯度裁剪防止 ViT 注意力层梯度爆炸 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() scheduler.step() # 每 10 个 epoch 验证一次 if (epoch 1) % 10 0: model.eval() psnr_list [] with torch.no_grad(): for hazy, clear in val_loader: hazy, clear hazy.cuda(), clear.cuda() pred model(hazy) psnr compute_psnr(pred, clear) psnr_list.append(psnr) print(fEpoch {epoch1}, PSNR: {sum(psnr_list)/len(psnr_list):.2f})验证时用 PSNR 和 SSIM 两个指标PSNR 看整体重建质量SSIM 看结构保持。在 NH-HAZE 上好的模型 PSNR 能到 20dB 以上SSIM 到 0.75 以上I-HAZE 和 O-HAZE 上 PSNR 能到 25dB 以上。5. 避坑与排查真实雾霾去雾训练中最容易翻车的 5 个点5.1 现象训练 loss 正常下降但验证 PSNR 不涨原因过拟合。I-HAZE 只有 30 对训练图ViT 参数量大很容易记住训练集。解决开强数据增强翻转旋转随机裁剪加 DropOut 或 DropPath把 embed_dim 从 64 降到 32或者用预训练权重初始化。5.2 现象输出图像整体偏暗或偏亮原因残差学习的符号搞反了。如果模型学的是清晰图减雾霾图但推理时用了加法就会偏。解决检查 forward 里是x - out还是x out残差学习的目标是雾霾层所以应该是减法。5.3 现象NH-HAZE 上 PSNR 比 I-HAZE 低 5dB 以上原因非均匀雾霾下全局注意力被高浓度区域主导低浓度区域欠拟合。解决在损失函数里加空间权重图对高浓度区域给更大权重或者用多尺度训练让模型同时看到不同浓度的雾霾。5.4 现象训练到一半 loss 突然变成 NaN原因ViT 的注意力 logits 数值过大softmax 后梯度爆炸。解决在 attention 里加attn attn - attn.max(dim-1, keepdimTrue)[0]做数值稳定或者把学习率降到 1e-4加梯度裁剪 max_norm0.5。5.5 现象测试时图像边缘出现明显拼接痕迹原因窗口注意力的窗口边界处信息不连续。解决测试时用 overlapping 窗口或者把 shift_size 设为 window_size//2让相邻窗口有重叠区域。6. 进阶技巧用 NTIRE2019 做跨数据集验证与模型部署6.1 跨数据集验证的正确姿势在 I-HAZE 上训完直接测 NH-HAZEPSNR 掉 3-5dB 是正常的。要缩小这个 gap我一般用两阶段训练先在 NTIRE2019 上预训练再在目标数据集上微调。微调时只解冻 decoder 和最后两个 stageencoder 的前两个 stage 冻结学习率降到 1e-5。# 冻结 encoder 前两个 stage for name, param in model.named_parameters(): if encoders.0 in name or encoders.1 in name: param.requires_grad False # 微调优化器只更新解冻的参数 optimizer AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr1e-5)6.2 推理加速与部署注意事项ViT 去雾模型在 1080p 图像上推理一次大概 200-500msV100要上生产环境需要做三件事第一用 TensorRT 或 ONNX Runtime 做图优化第二把模型量化到 FP16显存占用减半速度提升 30% 左右第三对于实时性要求高的场景把窗口大小从 8 降到 4注意力计算量降为原来的 1/4。# ONNX 导出示例 torch.onnx.export( model, torch.randn(1, 3, 256, 256).cuda(), dehaze_vit.onnx, opset_version13, input_names[input], output_names[output], dynamic_axes{input: {2: height, 3: width}} )导出时注意 dynamic_axes 要设对否则部署时只能跑固定尺寸。另外ONNX 对 LayerNorm 和 GELU 的支持在 opset 13 以上才完整低于这个版本会报错。6.3 一个我踩过的坑有一次在 O-HAZE 上训了一个模型PSNR 刷到 28dB结果部署到实际场景里遇到天空区域大面积过曝的雾霾图输出直接变成一片白。后来发现是训练集里没有高亮度雾霾样本模型对亮度分布过拟合了。解决办法是在数据增强里加了随机 gamma 变换模拟不同曝光条件下的雾霾。这个教训让我养成了一个习惯每次训完模型先拿几十张非数据集的真实雾霾图跑一遍看有没有系统性偏差再决定要不要上线。希望帮到你。本文还有配套的精品资源点击获取