
简介面向图像去雾算法研究与视觉Transformer技术的源码项目主要服务高校研究生、毕业设计团队以及有一定深度学习基础的算法工程师。算法以Uformer网络为核心覆盖数据切分、样本预处理、模型训练、效果验证和结果可视化等完整环节描述中提及的切图脚本负责将训练数据裁剪为256×256小图训练入口代码则集中管理批次大小、训练轮数、输入尺寸等关键参数便于使用者按设备条件灵活调整后重新训练。压缩包内共有338个文件约156.36MB以Python脚本为主辅以yaml配置文件、ipynb分析笔记、图片与动图结果展示、csv实验数据以及Markdown说明文档层次分明。实验数据中包括损失曲面与鲁棒性评测记录可帮助读者观察不同模型在受污染样本下的训练变化配套文档还写明了数据集存放位置和预处理步骤能够明显减少环境搭建与复现踩坑所花时间。目前已有291人学习浏览适合希望深入掌握视觉Transformer在图像去雾等低层视觉任务中落地细节的读者。1. 基于Vision Transformer的图像去雾为什么值得自己动手实现一遍清晨大雾里的监控画面、自动驾驶的前视摄像头、航拍图上的霾层——图像去雾一直是计算机视觉里“看起来简单、做起来翻车率高”的方向。传统暗通道先验在均匀雾天还能扛一到非均匀雾、浓淡交错的大场景就开始发灰、偏色、出现光晕。这两年Vision TransformerViT被引入去雾任务后效果提升明显核心原因只有一个ViT能把整张图的任意两个区域直接建立联系而雾的分布恰恰是全局性的。这篇笔记我会站在做项目的角度把“基于Vision Transformer的图像去雾算法”从原理、数据、训练、调参到落地避坑完整拆开给出可复现的PyTorch代码和参数配置。适合正在做CV课程设计、毕设或者刚入门深度学习算法想找一个“能讲清楚、能跑通、指标好看”的项目的人。2. 为什么是ViT图像去雾的全局依赖从哪里来2.1 大气散射模型与网络要学的东西图像去雾绝大多数工作建立在同一个物理模型上大气散射模型。一张有雾图像 I 可以看成清晰图像 J 经过透射率 t 衰减后再叠加大气光 A 的散射I(x) J(x) * t(x) A * (1 - t(x))t(x) 和景深有关通常写成 t(x) exp(-β * d(x))其中 β 是散射系数d(x) 是像素对应的场景深度。去雾任务本质上是已知 I反推 J。这里有两条路线一条是先用网络估计 t 和 A再套公式恢复 J物理约束强但估计误差会在除法里被放大另一条是端到端直接让网络输出 J简单粗暴但中间过程不可控。我一般在项目里选第一条路网络输出透射率图大气光用传统方法或一个轻量分支估计最后用公式恢复。这样做的好处是每一步都有可解释性论文里能画图答辩时也讲得清楚。网络要学的东西就从“预测一张图”变成了“预测一个 t 图t 的每个像素在 [0,1] 之间”任务难度明显下降。顺带一提如果你拿到的数据集本身有深度图还能用深度图直接合成干净且连续的 t比让网络盲猜稳定得多。2.2 ViT 的全局感受野如何匹配去雾任务CNN做去雾也不是不行但有个天然短板卷积的感受野是局部扩张的要靠堆深度才能让高层看到全局。遇到非均匀雾、大面积浓雾区浅层特征根本看不到远处像素的信息容易把一片区域整体处理错。ViT把图像切成patch每个patch变成一个token自注意力机制让任意两个token之间直接交互——第一层就能看到整张图。这正是去雾需要的远处浓雾和近处薄雾之间需要互相参照才能估计出合理的透射率梯度。用ViT去雾的代价也直接token数量随分辨率平方上涨256×256输入、patch_size8时序列长度是32×321024还不算太夸张要是上1024×1024的输入序列长度到16384显存直接爆炸。所以ViT去雾项目里patch_size和输入分辨率是两个决定成败的参数后面第四章会细说。2.3 网络主干PatchEmbed TransformerBlock 去雾头为了让你能在本地跑通我给出一个精简但结构完整的ViT去雾网络。它不追求刷榜胜在每一层都能讲明白改起来也方便。先看patch embeddingimport torch import torch.nn as nn class PatchEmbed(nn.Module): def __init__(self, in_ch3, patch_size8, dim512): super().__init__() self.proj nn.Conv2d(in_ch, dim, kernel_sizepatch_size, stridepatch_size) self.norm nn.LayerNorm(dim) def forward(self, x): # x: (B, 3, H, W) x self.proj(x) # (B, dim, H/p, W/p) B, C, Hp, Wp x.shape x x.flatten(2).transpose(1, 2) # (B, N, dim), N Hp*Wp return self.norm(x)这里用一个大步长卷积实现patch切分等价于先把图像切成不重叠的小块再分别做线性映射。LayerNorm放在embedding之后能稳定后续自注意力的输入分布。dim是最关键的超参代表每个token的向量长度论文里常见的是384或512再往上加对去雾这种中等复杂度任务收益就很有限了。然后是一个标准Transformer编码器块class TransformerBlock(nn.Module): def __init__(self, dim, heads8, mlp_ratio4, dropout0.1): super().__init__() self.norm1 nn.LayerNorm(dim) self.attn nn.MultiheadAttention( dim, heads, dropoutdropout, batch_firstTrue) self.norm2 nn.LayerNorm(dim) self.mlp nn.Sequential( nn.Linear(dim, dim * mlp_ratio), nn.GELU(), nn.Dropout(dropout), nn.Linear(dim * mlp_ratio, 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 xPyTorch内置的MultiheadAttention已经封装了Q、K、V的线性投影和注意力计算batch_firstTrue让输入输出形状都是(B, N, dim)省去手写矩阵乘法的麻烦。残差连接加Pre-LN先归一化再进注意力是ViT训练稳定的关键千万别为了“看上去更标准”改成Post-LN那会让训练前期loss下降慢得怀疑人生。最后把Transformer输出的token序列重新拼回图像形状并接一个透射率预测头class DehazeViT(nn.Module): def __init__(self, img_size256, patch_size8, dim512, depth6, heads8): super().__init__() n_patches (img_size // patch_size) ** 2 self.patch_embed PatchEmbed(3, patch_size, dim) self.pos_embed nn.Parameter( torch.randn(1, n_patches 1, dim) * 0.02) self.cls_token nn.Parameter( torch.randn(1, 1, dim) * 0.02) self.blocks nn.Sequential(*[ TransformerBlock(dim, heads) for _ in range(depth)]) # 透射率头每个token回归一块patch的透射率 self.t_head nn.Sequential( nn.Linear(dim, patch_size * patch_size), nn.Sigmoid()) self.patch_size patch_size self.img_size img_size def forward(self, x): B, _, H, W x.shape tokens self.patch_embed(x) tokens torch.cat([self.cls_token.expand(B, -1, -1), tokens], dim1) tokens tokens self.pos_embed tokens self.blocks(tokens)[:, 1:] # 去掉cls_token N tokens.shape[1] side int(N ** 0.5) tokens tokens.transpose(1, 2).reshape( B, self.patch_embed.proj.out_channels, side, side) t_map self.t_head(tokens.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) t_map nn.functional.interpolate( t_map, size(H, W), modebilinear, align_cornersFalse) return t_map这个去雾头是“教育版”设计每个token直接回归它对应patch大小的一块透射率再用双线性插值上采样回原分辨率。好处是结构简单、显存占用低缺点是patch边界上会有轻微不连续实际项目里可以用一个轻量的4层卷积头上采样到原图效果再提升一截。Sigmoid保证输出落在(0,1)正好对应透射率的物理范围。2.4 从原理到实现让公式变成可训练的结构网络输出的t_map并不能直接拿去算指标还需要结合大气光A把恢复图算出来。常见做法是用暗通道先验的思路估计A取有雾图像暗通道里最亮的前0.1%像素在这些位置上取原图的亮度均值。这个估计方法在绝大多数场景下够用而且是可微的可以直接放进训练图里。恢复环节的代码长这样def dehaze_with_model(model, haze): t model(haze) # (B, 1, H, W), 范围(0,1) t t.clamp(min0.1, max1.0) # 防止除零和过度增强 # 用暗通道前0.1%像素估计大气光 A dark haze.amin(dim1, keepdimTrue) k max(int(0.001 * dark.numel()), 1) flat dark.flatten(2) idx flat.topk(k, dim2).indices A torch.gather(haze.flatten(2), 2, idx.expand(-1, 3, -1)) A A.mean(dim2, keepdimTrue) # (B, 3, 1) J (haze - A) / t A return J, t, Aclamp的下限0.1是防除零的惯例值太小会让近景像素被放大得离谱太大又会让远景雾感残留。A的形状是(B, 3, 1)代表整张图共享一个全局大气光值对大多数户外场景是合理的近似如果场景里光源区域和雾区差异极大才需要考虑逐像素A的复杂模型。到这里一个完整的ViT去雾前向链路已经闭环接下来就是数据。3. 用PyTorch跑通最小去雾训练数据、损失与主循环3.1 训练数据从哪来SOTS子集与在线合成雾去雾训练最理想是拿到同一场景的“有雾/无雾”成对图像现实里这种数据很难拍。学术圈的标准做法是用合成数据RESIDE数据集里的SOTS子集提供了大量室内外成对样本是衡量去雾算法最常用的benchmark之一。如果你不方便下载还有一条完全自给自足的路用NYU-Depth的室内深度图套大气散射模型在线合成雾图。深度图给了公式里最关键的t(x) exp(-β*d(x))想合成什么浓度就合成什么浓度。两种数据的取舍SOTS是现成的图像质量和深度都是离线算好的省事但它的雾是固定的模型容易过拟合那一批beta参数。自合成的优势是beta和大气光A可以每个batch随机采样模型见过的雾场景更丰富真实场景迁移时不容易翻车。我一般偏爱后者训练前把深度图resize到256×256与RGB图对齐代码里实时合成。3.2 合成雾代码与beta采样策略import numpy as np def add_haze_np(clean, depth, beta_range(1.0, 2.5), A_range(0.7, 1.0)): 在线合成雾图: I J * exp(-beta*d) A * (1 - exp(-beta*d)) h, w clean.shape[:2] beta np.random.uniform(*beta_range) A np.random.uniform(*A_range, size(1, 1, 3)).astype(np.float32) t np.exp(-beta * depth.astype(np.float32) / 255.0) haze clean.astype(np.float32) * t A * (1 - t) return np.clip(haze, 0, 255).astype(np.uint8)depth要除以255归一化因为常见的深度图都存成0-255的灰度图不处理的话t会指数衰减到几乎为0合成出来的雾浓得看不见物体。beta_range的取值是个学问1.0-1.5近似薄雾2.0-2.5是浓雾。我一般按batch混合采样而不是全程固定一个范围让模型同时见过薄雾和浓雾避免训练集里“只有一种雾”导致测试时域差距过大。A的取值在0.7到1.0之间对应白色雾气的物理直觉。3.3 训练主循环与三件套损失有了数据和网络训练主循环最核心的是损失函数设计。只用L1损失恢复图容易发灰、边缘模糊只用SSIM图像锐度上来了但颜色容易偏。项目里我常用L1 SSIM的组合监督目标是“最终恢复图”让梯度同时穿过大气散射公式反传到t_predfor epoch in range(epochs): for haze, clean, depth in loader: haze, clean haze.cuda(), clean.cuda() t_pred model(haze).clamp(min0.1, max1.0) A estimate_atmosphere(haze) # 暗通道前0.1%像素均值 dehazed (haze - A) / t_pred A loss 1.0 * F.l1_loss(dehazed, clean) \ 0.3 * (1 - ssim(dehazed, clean)) opt.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step()ssim函数可以直接用PyTorch生态里的pytorch_msssim包一次前向返回一个标量用法和损失一样。为什么SSIM权重只有0.3因为它的梯度占比已经很大了权重超过0.5后loss会被结构相似度主导颜色精度反而不稳定。clip_grad_norm_在ViT训练里基本是必需品自注意力的梯度方差比CNN大不裁剪的话前几个epoch很容易出现loss突跳到NaN。3.4 超参配置表与显存预算给你一份我在单卡RTX 3090上跑过的保守配置直接抄就能动参数推荐值说明输入分辨率256×256再大优先用推理期tile不要硬撑训练patch_size8像素级任务别用16会丢细节dim5126层Transformer下性价比最高的值depth6显存紧张先减depth再减dimbatch_size83090 24G下无AMP能跑AMP可到16优化器AdamWlr1e-4, weight_decay1e-4学习率调度Cosine Annealing100 epochs衰减到1e-5混合精度torch.cuda.amp显存减半训练提速约30%epoch100第40轮后指标才明显爬升别早停这张表是“能跑通”的基线。如果你的卡是16G甚至12G先把batch_size降到4用梯度累积补足等效batch还撑不住就把depth从6降到4效果损失远小于把patch_size改成16。ViT去雾对patch_size最敏感这是我在多个数据集上反复踩出来的结论。4. 参数调优从“能跑”到“看得过去”的4个关键点4.1 patch_size、depth、dim三者的取舍逻辑ViT去雾里最容易被忽略、影响却最大的超参是patch_size。patch_size4时token数多模型能看到更细的纹理边界但自注意力计算量是patch_size8的4倍patch_size16时计算省了但每个token覆盖区域太大去雾结果会出现明显的块状伪影尤其是雾的浓淡分界线附近。我的一般选型路径是配置token数256²输入相对计算量适用场景patch4, dim256, depth44096高小图、离线精修patch8, dim512, depth61024中通用项目基线推荐patch16, dim768, depth8256低大图粗去雾不追求边界记住一个原则当显存成为瓶颈时优先削depth其次削dim最后才动patch_size。depth从6减到4感受野能力下降不明显dim从512减到256信息瓶颈立刻出现恢复图的颜色细节会变糊patch_size从8改到16则是直接丢空间分辨率伪影问题无解。4.2 学习率与优化器Vision Transformer的隐藏脾性ViT对学习率比CNN敏感得多。CNN你用SGD0.01也能跌跌撞撞收敛ViT用SGD经常前期就卡死。我所有ViT去雾项目都只用AdamW初始学习率锁在1e-4到2e-4之间配Cosine Annealing。有个容易翻车的细节ViT里position embedding和cls_token的初始化幅度要小0.02量级而最后的透射率头可以单独给稍大的学习率或初始化倍数否则网络前几十个batch会一直在“忽略输入、输出接近常数”的状态里挣扎。还有个来自实践的经验前10个epoch用warmup从1e-5线性升到1e-4。这10个epoch看起来loss没怎么降但在为自注意力的Q/K投影积累稳定的梯度统计量。直接上大学习率前几个batch就可能把patch embedding的特征分布打散后面很难救回来。4.3 损失函数权重分值与人眼之间的平衡项目标题挂在“高分项目”上这就避不开一个问题评测指标分数高不代表人眼看过去舒服。PSNR对整体亮度偏移非常敏感模型稍微把图调亮一点PSNR涨很多画面却发灰SSIM对结构敏感但对颜色不够敏感。训练时只压L1会得到高PSNR但视觉发灰的图只压SSIM会得到边缘锐利但偏色的图。我的做法是L1权重1.0 SSIM权重0.2~0.3如果数据量超过一万张再加一项VGG感知损失权重0.05起步。感知损失的作用是让恢复图的纹理和清晰图在特征空间里一致能明显拉高主观质量但它梯度来自预训练VGG对小尺寸训练图有偏置数据量不足时反而会让颜色怪掉。先在验证集上算PSNR和SSIM再挑三五张典型图肉眼过一遍比只盯一个指标靠谱得多。4.4 训练节奏与检查点策略ViT去雾的收敛是“先粗后细”前30个epoch透射率图的大致轮廓出来恢复图已经能看到60个epoch边界和细节才开始稳定100个epoch后PSNR的提升可能只有0.2dB但视觉上雾感会干净很多。我在第20、40、60、80个epoch都会保存检查点并在验证集上计算PSNR/SSIM。这样有个后悔药万一第90轮开始过拟合还能退回第80轮的权重而不是重训一遍。不要指望只看训练loss决定收敛。训练loss里包含SSIM项它会在某个点后平缓下降但肉眼效果仍在提升。真正可靠的信号是验证集PSNR和几张固定测试图的定性对比。固定测试图很重要——每一次验证都用同一批图你才能看出来模型是在稳定改善还是随机抖动。5. 图像去雾落地避坑5条血泪经验5.1 推理尺寸与训练不一致网格状伪影现象训练用256×256推理时直接喂1024×1024的大图结果输出图像表面出现规则的网格纹理雾的边界像被切成了很多小方块。原因ViT的位置编码是训练时按固定token数学出来的。推理分辨率变大后token数变多位置编码和输入序列长度对不上即便很多实现做了插值插值后的位置编码分布和训练时的分布有偏差自注意力就会受到周期性的位置扰动。解决两种做法二选一。要么推理时把大图切成256×256的patch分块处理再拼接边缘平滑过渡要么在forward里对pos_embed做双线性插值到目标长度并在插值后做短距离的微调训练。项目里图省事我用前者多写一个tile函数而已效果稳定且不需要改模型。5.2 合成雾域差距beta范围太窄真实雾天翻车现象SOTS或自合成的数据上PSNR能到28dB以上把模型拿到真实雾天照片上恢复图要么雾没去干净要么局部过增强变成一块黑斑。原因合成雾用的beta和A都是理想值真实雾天还有大气湍流、非均匀散射、传感器噪声数据分布对不上。如果你训练时beta固定为2.0这一档模型只学会了“去中等浓度的雾”其他浓度全靠泛化硬撑。解决训练时把beta_range放宽到(0.8, 3.0)并且按场景分桶采样——每批里薄雾、中雾、浓雾各三分之一。大气光A也不要固定为1.0在0.6到1.0之间随机。这个小改造能让真实雾天测试的稳定性提升一大截是性价比最高的数据策略。5.3 透射率头激活函数选错恢复图发黑现象模型训练到中途输出图像突然大面积变暗或发黑透射率图看起来像马赛克且loss在某个区间不停震荡。原因透射率头的输出原本设计成Sigmoid有人为了“加速收敛”换成ReLU或直接线性输出。透射率超过1.0或小于0之后恢复公式里的除法会把像素推到极端值整体图像偏暗甚至变黑。解决head末尾保持Sigmoid并在前向强制clamp到[0.05, 1.0]区间。如果发现训练中期暗块反复出现检查一下是不是head输出的均值漂移了可以给Sigmoid前一层加个LayerNorm输出分布会更稳定。这一条是我自己翻车翻出来的换成Linear输出后白白浪费了三天训练时间。5.4 PSNR高分但图像发灰只有L1的代价现象验证集PSNR 29dBSSIM也到了0.92但肉眼看恢复图像蒙了一层灰纱对比度明显偏低树叶和天空交界处发闷。原因L1损失做的是逐像素绝对值平均模型最安全的选择是输出接近训练集平均值的图这样每个像素误差都小但图像的对比度和色彩饱和度都被压低了。PSNR被这种“保守策略”拉高却不符合人眼对“清晰”的感知。解决在训练损失里加入SSIM它本质是结构和亮度对比的度量能逼着模型把对比度拉回来。如果SSIM加了还是发灰再检查一下是不是输入图像归一化时用了错误的均值标准差导致模型看到的输入整体偏暗。5.5 显存OOMpatch_size8直接爆显存现象按本章的基线配置训练10G显存的卡在第二个epoch就OOM报错信息指向MultiheadAttention的前向。原因自注意力的显存占用是token数的平方。256×256输入、patch_size8时单张图的注意力矩阵是1024×1024batch_size8就同时持有8个这样的矩阵显存撑不住。解决把batch_size降到4开启torch.cuda.amp混合精度再用梯度累积模拟batch_size16的效果。如果还想加大模型的表达能力就维持patch_size8不动把depth降到4。记住显存紧张时先降batch再降depthpatch_size尽量保持8这个优先级能保证去雾质量的崩坏速度最慢。6. 进阶从验证到轻量化的三个技巧模型训完接下来要解决的是“怎么让别人相信它有用”和“怎么让它跑得快”。验证环节有个很容易踩的细节计算PSNR前必须先把预测图和GT转成相同的数值范围并且只在RGB空间计算别在BGR上算。更严谨一点我会把图像裁掉靠近边缘的8个像素再算指标因为去雾网络在边缘存在padding效应把边缘算进去会低估真实质量。SSIM方面我建议同时报告SSIM和UQI后者对亮度偏移更敏感两个指标一起看能判断模型是“真去雾了”还是“只是提亮了”。轻量化方向上常用做法是蒸馏用训练好的DehazeViTdepth8当teacher训一个depth4的小学生网络蒸馏损失用teacher输出的透射率图做软标签加上L1。这样能把模型体积砍掉近一半推理速度提升可观而PSNR损失通常控制在0.3dB以内。如果再配合剪枝算法把注意力头数从8剪到6精度损失更小适合放到边缘设备上。我自己做这个项目的最深感受是ViT去雾的上限确实比CNN高但代价是训练成本和调参复杂度都上了一个台阶。一开始我总想着把模型做大结果卡在显存和过拟合里出不来后来老老实实从patch_size8、depth6的基线出发一步步加数据和损失效果反而稳定。如果你也要用这个标题交作业或发论文切记先把训练闭环跑通再去动网络结构。希望帮到你。本文还有配套的精品资源点击获取