简介一套基于Vision Transformer的图像去雾算法研究与实现资料包内含源码与文档说明适合计算机视觉方向研究生、算法工程师以及图像复原爱好者用于学习与二次开发。项目采用Uformer架构覆盖从NH-HAZE数据集预处理到模型训练的全流程包含将训练图像切分为256×256小图的generate_patches_SIDD.py预处理脚本以及支持--arch Uformer、--nepoch 270、--batch_size 32等参数配置的My_train.py训练代码并配套YAML配置文件、CSV损失记录与结果图可帮助复现去雾实验并分析模型收敛情况。压缩包共338个文件包含204个Python脚本、16个YAML配置、12个CSV记录、9个IPython笔记及8个Markdown文档另有PNG/GIF结果图与TXT日志整体约156MB目录结构清晰。目前已有291人学习下载。配套文档说明详细解读Vision Transformer在图像去雾中的实现思路、数据组织方式及训练调参要点适合作为毕业设计、课程项目或科研工作的参考基线。1. 图像去雾为什么从卷积转向 Transformer一次视觉任务的分水岭雾天拍摄的监控画面和自动驾驶感知图像往往整体发灰、对比度极低这类图像的共同特征是清晰度损失并不只发生在局部而是整幅图像都被大气散射压暗。传统去雾算法如暗通道先验在天空区域和白色物体上经常翻车原因是它假设“每个局部patch里总有一个通道的强度趋近于零”这个假设在亮度均匀的区域根本不成立。而基于 Vision Transformer 的图像去雾算法把整幅雾图切成一串视觉词元patch通过自注意力机制建立远距离像素之间的联系正好补上了卷积局部感受野的短板。本文会从大气散射模型讲起逐步给出可运行的最小实现、参数配置、评估方法和踩坑记录适合正在做图像去雾课题或用 Transformer 做底层视觉任务的同学照着复现。2. 从大气散射模型到 Vision Transformer去雾任务为什么吃这套2.1 大气散射模型与 ViT 的对应关系图像去雾的理论基础是大气散射模型Atmospheric Scattering ModelI(x) J(x) * t(x) A * (1 - t(x))其中 I(x) 是观测到的雾图J(x) 是我们要恢复的清晰图像t(x) 是透射率A 是全局大气光。这个式子展开来看去雾任务的核心就是估计两个未知量逐像素的透射率 t(x) 和全局大气光 A。一旦这两项估计出来清晰图像可以通过 J(x) (I(x) - A) / t(x) A 直接算出。关键细节在于透射率 t(x) 的空间分布。在同一个场景里近处的物体透射率高、远处的山体透射率低t(x) 的变化往往跨越整个图像平面。卷积神经网络受限于卷积核的大小通常只能看到 3x3 或 7x7 的局部邻域要让它学到“远山和近景之间的透射率渐变关系”必须堆很多层才能把感受野扩到足够大。Vision Transformer 天然没有这个问题——它把图像分成 patch 之后每个 patch 的 token 会与全图所有其他 patch 的 token 做注意力计算一步就建立了全局依赖。这意味着在估计透射率时模型可以直接比较图像左上角和右下角的亮度差异从而判断深度梯度。2.2 为什么选择 Vision Transformer 而不是堆更大的卷积核做去雾之前我也试过把普通卷积网络加宽加深比如把 UNet 的 encoder 部分替换成 ResNet-50效果在均匀雾天情况下还行但遇到浓雾和薄雾同时出现的场景就明显不够用。深层卷积网络在训练时还有个常见问题梯度在跨层传播过程中会逐渐消失深层特征容易丢失浅层的细节纹理。Transformer 的残差连接和 LayerNorm 机制在一定程度上缓解了这个问题更重要的是它的注意力头可以并行地关注不同的空间关系——某些头关注全局亮度分布某些头关注边缘细节这种多样化的表征能力恰好对应去雾问题中“估计透射率需要全局信息、恢复纹理需要局部信息”的双重需求。另外值得注意的一点是Vision Transformer 在图像分类任务上的成功并不能直接迁移到去雾任务上来因为去雾是稠密预测任务输出要求与输入同分辨率的逐像素结果。学术界通常的做法是在 ViT 后端接入一个解码器decoder结构或者直接把 Transformer 作为特征提取器再配合卷积上采样层把特征图恢复到原始尺寸。本文的实现就是走这条路ViT 负责全局建模后接轻量级卷积解码器恢复细节。3. 基于 Vision Transformer 的图像去雾最小可跑实现从数据到模型逐段拆解3.1 数据准备先造雾再学去雾公开去雾数据集如 RESIDE的下载和预处理需要较大的网络成本日常实验调试阶段最常用的做法是先用清晰图像合成雾图。合成公式严格按大气散射模型来雾图 清晰图 x 透射率图 大气光 x (1 - 透射率图)透射率图通常用一个随深度变化的平滑张量来表示深度信息可以用随机生成的分层梯度图近似。import torch import torch.nn.functional as F import numpy as np def generate_haze(clean_image, min_t0.3, max_t1.0, airlight0.8): 按大气散射模型合成雾图 clean_image: (B, C, H, W)取值范围 0~1 的清晰图像 min_t: 透射率下限越小代表雾越浓 airlight: 大气光强度0~1 之间越高整体越偏灰白 batch_size, _, h, w clean_image.shape # 生成平滑的透射率图先随机生成粗略深度图再做高斯模糊 depth torch.rand(batch_size, 1, h // 8, w // 8) depth F.interpolate(depth, size(h, w), modebilinear, align_cornersFalse) depth F.avg_pool2d(depth, kernel_size31, stride1, padding15) # 平滑 # 把深度归一化到 [min_t, max_t] 区间 t min_t (max_t - min_t) * (1.0 - depth) # 大气光按全局常数处理真实场景多为 (0.7~0.9) 的灰度值 a torch.full_like(clean_image, airlight) haze clean_image * t a * (1.0 - t) return haze, t这段代码里最难调的参数是 min_t。取值越小雾越浓训练出的模型对浓雾的适应能力越强但训练难度也越大因为透射率过低时清晰图像的信号衰减严重模型容易把输出直接推向“全图抹平”。我一般调试阶段用 min_t0.3跑通流程后用 0.2 加强鲁棒性。avg_pool2d 的 kernel 大小决定透射率图的平滑程度kernel 越大深度渐变越自然但如果设得太大近景和远景之间就没有层次感了。3.2 Patch Embedding把雾图切成视觉词元Vision Transformer 的第一步是把输入图像切分成固定大小的 patch然后通过线性映射把每个 patch 压成一个向量。这一步和 NLP 里把句子切成词元再映射成 embedding 是一个道理只是在图像里我们用的是二维 patch 而非一维词元。import torch import torch.nn as nn class PatchEmbed(nn.Module): def __init__(self, in_channels3, patch_size4, embed_dim256): super().__init__() self.proj nn.Conv2d(in_channels, embed_dim, kernel_sizepatch_size, stridepatch_size) def forward(self, x): x: (B, 3, H, W) 输出: (B, N, embed_dim)N H/4 * W/4 x self.proj(x) # (B, embed_dim, H/4, W/4) b, c, h, w x.shape x x.flatten(2) # (B, embed_dim, N) x x.transpose(1, 2) # (B, N, embed_dim) return xpatch_size 的选取直接影响计算量和感知粒度。patch4 时一张 256x256 的输入会被切成 64x644096 个 patchTransformer 的注意力计算量是 O(N^2)N4096 时自注意力矩阵是 4096x4096显存压力会显着上升。patch8 时 N1024计算量小了四倍但每个 patch 内部的信息混合更粗边缘细节的恢复能力变弱。实验对比下来256x256 输入配 patch4 是效果和显存的平衡点如果你的 GPU 显存只有 8G 左右建议把输入降到 192x192 或直接用 patch8。3.3 Transformer Encoder 加透射率预测头Transformer Encoder 是整条网络的核心由多头自注意力层和前馈网络交替堆叠而成。和 ViT 分类模型不同的是去雾任务的输出需要保留空间结构所以我要在 encoder 后面再接一个轻量卷积解码器把特征图上采样回原分辨率。import torch import torch.nn as nn import math class TransformerBlock(nn.Module): def __init__(self, embed_dim, num_heads, mlp_ratio4.0, dropout0.1): super().__init__() self.norm1 nn.LayerNorm(embed_dim) self.attn nn.MultiheadAttention(embed_dim, num_heads, dropoutdropout, batch_firstTrue) self.norm2 nn.LayerNorm(embed_dim) self.mlp nn.Sequential( nn.Linear(embed_dim, int(embed_dim * mlp_ratio)), nn.GELU(), nn.Dropout(dropout), nn.Linear(int(embed_dim * mlp_ratio), embed_dim), nn.Dropout(dropout), ) def forward(self, x): x x self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0] x x self.mlp(self.norm2(x)) return x class ViTDehaze(nn.Module): def __init__(self, img_size256, patch_size4, embed_dim256, num_heads8, depth6): super().__init__() self.patch_embed PatchEmbed(in_channels3, patch_sizepatch_size, embed_dimembed_dim) self.cls_token nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed nn.Parameter(torch.randn(1, (img_size // patch_size) ** 2 1, embed_dim)) self.blocks nn.Sequential(*[ TransformerBlock(embed_dim, num_heads) for _ in range(depth) ]) self.norm nn.LayerNorm(embed_dim) # 解码器从序列恢复到图像尺寸 self.decoder nn.Sequential( nn.Conv2d(embed_dim, embed_dim // 2, kernel_size3, padding1), nn.GELU(), nn.Upsample(scale_factor2, modebilinear, align_cornersFalse), nn.Conv2d(embed_dim // 2, embed_dim // 4, kernel_size3, padding1), nn.GELU(), nn.Upsample(scale_factor2, modebilinear, align_cornersFalse), nn.Conv2d(embed_dim // 4, 3, kernel_size3, padding1), ) def forward(self, x): B x.shape[0] tokens self.patch_embed(x) # (B, N, embed_dim) # 前置 cls_token 和位置编码 cls_tokens self.cls_token.expand(B, -1, -1) tokens torch.cat([cls_tokens, tokens], dim1) tokens tokens self.pos_embed tokens self.blocks(tokens) tokens self.norm(tokens) # 去掉 cls_token恢复成特征图并送入解码器 tokens tokens[:, 1:, :] feat_h feat_w int(math.sqrt(tokens.shape[1])) feat tokens.transpose(1, 2).reshape(B, -1, feat_h, feat_w) out self.decoder(feat) return out代码里两个容易忽略的细节一是 cls_token 初值用了零初始化位置编码用了 randn 而不是固定正弦编码因为图像 patch 的空间关系比较复杂随机位置编码配合训练中学习比手工设计的三角函数编码更灵活。二是 decoder 的上采样直接用 bilinear没用转置卷积。转置卷积在 stride2 时容易产生棋盘伪影bilinear 插值虽然理论上更“朴素”但在去雾这种需要平滑输出的任务里反而能减少高频噪声的引入。3.4 训练脚本与损失函数组合去雾任务的损失函数不能只用 L1 或 MSE。L1 损失对整体亮度恢复比较稳健但容易让输出图像纹理过于平滑MSE 损失对异常值敏感。业界最常用的做法是 L1 损失 SSIM 损失的组合有条件的话再加上感知损失Perceptual Loss。import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader def ssim_loss(pred, target): 简化版 SSIM 损失window 大小 11 c1 (0.01 * 1.0) ** 2 c2 (0.03 * 1.0) ** 2 kernel torch.ones(1, 1, 11, 11) / 121.0 kernel kernel.to(pred.device) mu_x F.conv2d(pred, kernel, padding5) mu_y F.conv2d(target, kernel, padding5) mu_xx F.conv2d(pred * pred, kernel, padding5) mu_yy F.conv2d(target * target, kernel, padding5) mu_xy F.conv2d(pred * target, kernel, padding5) sigma_x mu_xx - mu_x * mu_x sigma_y mu_yy - mu_y * mu_y sigma_xy mu_xy - mu_x * mu_y ssim_map ((2 * mu_x * mu_y c1) * (2 * sigma_xy c2)) / \ ((mu_x * mu_x mu_y * mu_y c1) * (sigma_x sigma_y c2)) return 1.0 - ssim_map.mean() def dehaze_loss(pred, target, alpha0.8): l1 F.l1_loss(pred, target) ssim ssim_loss(pred, target) return alpha * l1 (1 - alpha) * ssim # 训练循环关键片段 model ViTDehaze(img_size256, patch_size4, embed_dim256, num_heads8, depth6) optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max50) for epoch in range(50): for batch in dataloader: haze_img batch[haze].cuda() clean_img batch[clean].cuda() pred model(haze_img) loss dehaze_loss(pred, clean_img) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() scheduler.step()alpha 参数控制 L1 和 SSIM 的权重。alpha 越大训练过程越稳定早期收敛越快但图像边缘容易发糊alpha 调小后 SSIM 主导纹理保持更清晰但在雾很浓的区域训练容易波动。我试过不同组合后alpha0.8 是省心的默认值先让它收敛到比较低的 loss再隔 20 个 epoch 把 alpha 降到 0.6 做微调效果会比固定权重好一些。clip_grad_norm_ 这行不能省Transformer 的梯度范数波动远大于卷积网络不裁剪梯度的话偶尔一个 loss 突刺就能让训练发散。4. 训练配置与评估PSNR、SSIM 之外还要看什么4.1 一组能直接用的超参数超参数组合决定了训练能否收敛以及收敛后效果的上限下面这组参数是我在单张 12G 显存显卡上调试出来的平衡点。参数推荐值调整方向输入尺寸256x256显存不够时降到 192但 patch_size 要同步保持整除关系patch_size4越大速度越快但细节恢复越差最大不建议超过 8embed_dim256调大能提升表征能力多头注意力的显存占用会指数上升num_heads8需要让 embed_dim 能被 num_heads 整除transformer depth6数据集小时 depth 超 8 容易出现过拟合batch_size8单卡 12G 上限超过 8 容易 OOM初始学习率1e-4AdamW 配 1e-4 比较稳比 3e-4 更不容易出现 loss 震荡训练 epoch50 20 微调前 50 个 epoch 用 alpha0.8后 20 个 epoch 用 alpha0.6这里有个容易误判的点不是 embed_dim 越大越好。embed_dim 从 256 调到 512 后PSNR 提升通常只有 0.3~0.5dB但显存占用几乎翻倍训练时间也肉眼可见地变长。如果数据集规模不大几千张图embed_dim256 已经能覆盖绝大多数去雾场景的表达需求。4.2 评估流程先看量化指标再做人眼判断去雾任务最常用的两个指标是 PSNR峰值信噪比和 SSIM结构相似性。PSNR 衡量像素级差异值越高越好SSIM 衡量结构一致性值越接近 1 越好。这两个指标在完全合成雾的数据上表现稳定但做对比实验时只报这两个数是不够的因为它们对“整体亮度偏移”非常敏感——把输出整体调暗 5%PSNR 会掉一大截但人眼看起来可能并不觉得多难受。import numpy as np from skimage.metrics import peak_signal_noise_ratio, structural_similarity def eval_dehaze(pred, target): # pred 和 target 都是 0~1 的 float 数组HWC 格式 pred np.clip(pred, 0.0, 1.0) target np.clip(target, 0.0, 1.0) psnr peak_signal_noise_ratio(target, pred, data_range1.0) ssim structural_similarity( target, pred, data_range1.0, channel_axis-1 ) return psnr, ssim评估时我一般跑三组对比原始雾图直接算指标作为 baseline、暗通道先验算法作为传统方法代表、ViT 去雾模型。这也是一篇论文或一份课程设计报告里最基础的三行对比表。但建议额外统计一个“天空区域占比”指标对雾图按亮度做阈值分割统计高亮区域在输出中的细节保留程度。很多时候整体 PSNR 很好看天空区域却被处理成一块平板视觉上非常出戏。4.3 视觉检查清单评估环节除了跑指标我习惯固定抽 10 张不同场景的图做肉眼检查重点看三个位置一是天空和物体交界处有没有白边二是远处低对比度区域有没有出现色块或噪声三是近景纹理有没有被过度锐化导致边缘发黑。这三个位置是去雾模型最容易出问题的地方量化指标很难完全反映。5. 避坑ViT 去雾最容易翻车的四个位置5.1 透射率越界导致输出黑斑现象训练到中途验证集上偶尔会出现全黑或带大片黑斑的输出图单张图 loss 正常但整体训练曲线开始抖动。原因模型预测的透射率没有做数值约束直接进入 J(x) (I(x) - A) / t(x) 的反演公式当 t(x) 预测值接近 0 时除法会把噪声无限放大得到接近全黑的图像。Transformer 在训练初期特别容易在某些 patch 上输出极小的透射率值。解决在解码器最后一层卷积后面加一个 sigmoid 激活把透射率输出限制在 (0, 1) 区间并且在使用时对透射率做下限截断例如 t max(t, 0.1)。这比在 loss 里加透射率平滑项更直接有效。5.2 patch 边界伪影现象输出图像能隐约看到网格状痕迹尤其在 256x256 输入、patch_size8 时特别明显放大后像马赛克。原因Patch Embedding 用 stridepatch_size 的卷积做切分相邻 patch 之间没有重叠区域Decoder 在上采样恢复时每个 patch 内部的特征独立性太强patch 之间缺少平滑过渡。解决第一个方法是把 patch_size 从 8 降到 4网格痕迹会明显减轻第二个方法是在 Decoder 的第一层卷积改用 kernel_size5、padding2让卷积核有机会跨 patch 混合信息第三个更彻底的做法是在 Patch Embedding 前加一层 stride1、kernel_size3 的预处理卷积把相邻像素的信息预先融合。5.3 小数据集上训练发散现象训练集只有几百张合成雾图时前几个 epoch 的 loss 在 0.1~0.3 之间反复跳怎么调学习率都压不住甚至出现 loss 越训越高的反直觉情况。原因Vision Transformer 的参数空间比同体量卷积网络更大对数据量的要求相应更高。几百张图不够喂饱 6 层 Transformer 的注意力权重模型容易把一些 patch 的注意力分数学到极端位置直接记住训练集噪声。解决我在这种情况下的首选是加数据增强把训练集做随机裁剪尺寸范围 0.8~1.0、水平翻转、随机亮度扰动乘子 0.85~1.15。注意不要在增强里加高斯噪声因为去雾任务本身就要从低信噪比的雾图中恢复信号额外噪声会让任务更难。5.4 评估指标和视觉感受脱节现象PSNR 均值 28dB、SSIM 0.92数字很漂亮但放大单张输出图时发现天空部分有肉眼可见的带状条纹columns 方向上颜色不连续。原因PSNR 是逐像素差异的均值它的惩罚机制对高频噪声不算敏感——一条细带子只占整图的百分之一像素不足以显著拉低 PSNR。而 SSIM 在局部窗口内计算结构相似度带状条纹如果宽度小于窗口尺寸同样很难被发现。解决在评估脚本里加一个高频分量统计把输出图像减去高斯模糊版本得到高频残差图统计残差图的能量分布。如果个别局部区域的高频能量远高于全图中位数基本可以断定有伪影或带状条纹。这个指标虽然不能写进实验对比表但对排查模型问题很有用处。6. 进阶用引导滤波把 ViT 输出边缘救回来ViT 去雾模型在整体亮度和色彩恢复上表现不错但输出图像的边缘细节有时不如基于卷积的传统方法锐利。这是注意力机制的固有取舍自注意力擅长捕捉全局依赖对高频局部结构的响应天然弱于卷积。一个成熟的工程做法是把 ViT 的输出和经典引导滤波结合起来——把 ViT 预测的透射率图作为引导图对去雾结果做边缘保持的滤波处理。import cv2 import numpy as np def guided_filter_dehaze(dehazed_img, guidance_img, radius16, eps1e-3): 用透射率图为引导对去雾输出做边缘保持平滑 dehazed_img: (H, W, 3) float 0~1 guidance_img: 单通道灰度图常用透射率估计 t 或雾图亮度 # 转成 float32 满足 OpenCV 输入要求 guidance guidance_img.astype(np.float32) dehazed dehazed_img.astype(np.float32) # 对每个颜色通道做引导滤波 result np.zeros_like(dehazed) for c in range(3): result[:, :, c] cv2.ximgproc.guidedFilter( guidance, dehazed[:, :, c], radius, eps ) return np.clip(result, 0.0, 1.0)radius 和 eps 这两个参数直接影响平滑程度。radius 太小小于 8时滤波窗口覆盖不了纹理区域边缘保留效果不明显radius 太大大于 32会把真正的边缘也磨平。eps 是正则项控制“什么程度的差异算边缘”——取值越小滤波越盲目相信引导图的边缘结构但引导图透射率图本身的误差也会被放大取值过大则退化成普通均值滤波。我通常从 radius16、eps1e-3 起步量化指标不会改善太多人眼观感上边缘更利落。如果想要更极致的边缘保持效果可以把引导图换成原雾图的亮度通道。原理是原雾图在有纹理的位置天然具备高频信息这些位置也正好是恢复清晰图像时最该保留细节的位置。最后说一个我的习惯每次训练完不急着收工把五六张典型测试图单独存下来分别用轻雾、中雾、浓雾三种强度合成看模型在三种退化程度下的表现差异。经常出现的情况是轻雾场景下 ViT 去雾不如简单直方图均衡因为轻雾图的对比度损失本身很小复杂模型的过度处理反而引入伪影。如果遇到这种现象可以把轻雾样本单独筛出来用更小的损失权重参与训练或者干脆在推理时对预测透射率做一个置信度阈值判断。希望这一篇能帮你在 ViT 去雾这条路上少走一些我走过的弯路。本文还有配套的精品资源点击获取