简介面向图像去雾与视觉Transformer研究者和开发者提供基于Vision Transformer算法的完整Python实现含源代码、预训练权重与详细使用说明。压缩包共340个文件约156MB其中204个Python脚本负责模型构建、训练参数配置与推理流程16个YAML文件定义实验环境与超参数12个CSV记录训练日志及评价指标9个Jupyter Notebook提供交互式演示另有GIF对比图与效果图辅助理解方便直观比较算法性能。已有1446人浏览学习。资源附带了按不同数据集划分的预训练权重My_best_model文件夹可通过--pretrain_weights参数指定权重路径继续训练训练所需的补丁大小--train_ps等核心参数集中在option.py中设置便于复现实验。此外项目介绍文档和详细使用说明覆盖了从环境配置到模型评估的完整链路适合有一定深度学习基础、想深入理解ViT在图像去雾低层视觉任务中应用的开发者。1. 为什么去雾任务越来越多人转向Vision Transformer而非纯卷积网络监控视频、无人机航拍、自动驾驶感知遇到雾天基本全靠后期算法救。过去几年暗通道先验和CNN去雾模型是主流但做过落地的人都有一个共同体验模型在合成雾评测集上PSNR能到30dB以上一到真实雾图就翻车。这个现象背后的原因很反直觉——CNN引以为傲的局部归纳偏置恰恰在去雾任务上成了包袱。雾的影响是全局的同一张图里近景清晰、远景发白透射率分布跨越整个空间Vision Transformer通过全局自注意力把每个图像块和全图所有块建立联系天然适合拟合这种长程、浓度渐变的物理过程。这篇笔记要讲清楚基于Vision Transformer的图像去雾算法怎么做Python工程从数据准备、模型搭建到训练推理怎么跑通以及那些不跑一遍根本发现不了的坑。2. 从大气散射模型到全局自注意力去雾任务的原理与选型理由2.1 大气散射模型去雾问题到底在解什么去雾的物理基础是大气散射模型写成公式是 I(x) J(x)t(x) A(1 - t(x))。I是相机收到的模糊图J是要恢复的干净图t是透射率0到1之间A是全局大气光。雾越浓t越小远端像素的贡献被A压得越厉害。算法要么估计t和A再逆推J要么端到端直接回归J。问题在于一张模糊图对应无穷多组(J, t, A)这是个病态反问题模型必须靠数据和先验把解空间压住否则恢复结果会偏色或者丢失细节。早年方法用先验约束t比如暗通道先验假设局部小块里至少一个颜色通道接近0最初的去雾效果确实惊艳但遇到天空、白色物体、明亮灯光这些违反先验的区域就露馅。这类方法的输出常常带一圈黑边或者把白色墙壁恢复成灰紫色。根本原因在于它用局部统计量估计全局物理量天然缺少跨区域推理能力。CNN方法好一些但受感受野限制想要建模远处浓雾和近景清晰的连续分布需要很深的网络不断膨胀感受野训练难度和推理成本都上升而且理论感受野不等于有效感受野——深层卷积真正依赖的区域往往远小于理论值远距离像素之间的关联仍然很弱。2.2 ViT的自注意力凭什么能拟合透射率分布Vision Transformer把图像切成固定大小的patch每个patch拉平成token经过一层线性投影后加上位置编码送进Transformer Encoder。核心是自注意力每一层里每个token都会和全图所有token计算相关性输出权重加权融合。这意味着第一层就能看到整张图的上下文远处天空的信息可以直接影响近景地面块的估计不需要靠层数堆叠。对去雾任务这个性质非常关键透射率图通常是全局平滑、局部突变的混合结构全局注意力负责捕捉平滑渐变前馈网络负责拟合局部突变两类信息在一个模块里共同决策。具体看Query、Key、Value的角色某个patch的Query和其他所有patch的Key做点积得到注意力权重再去加权对应的Value。响应值越高的token对当前patch的恢复结果影响越大。模型训练的过程就是学会让远处浓度相似的patch互相增强、让边界处反差大的patch互相抑制。位置编码还保留了patch之间的空间关系不会因为全局注意力丢掉图像结构。与分类任务不同去雾不用CLS token做全局表征而是保留每个token的输出reshape回二维特征图后再逐像素回归这样输出天然对齐输入位置不需要额外的对齐模块。2.3 与CNN、基于窗口注意力的Transformer变体对比什么时候选ViT不至于踩坑CNN在小雾、局部雾场景下已经够用且推理快、显存低但感受野是硬伤。基于窗口注意力的Transformer变体用窗口注意力把自注意力限制在局部窗口内计算量降了但跨窗口信息需要靠移位窗口来交换实现复杂度和调参难度都上升。纯ViT在这个任务上的优势是趁手它吃全局建模patch_size和层数可以直接参考成熟分类配置上手路径清晰代价是长序列注意力显存占用高训练收敛比CNN慢。方案全局建模能力显存占用实现复杂度适用场景CNN弱感受野有限低低轻度雾、实时推理窗口注意力变体中靠跨窗口交换中高高分辨率图、资源受限纯ViT强第一层即全局高中均匀雾、追求画质上限我的选择习惯是白天均匀雾、数据量大用ViT做主干端到端回归夜间或非均匀雾ViT结构保留但要额外加物理模型先验分支。从复现一个开源Python工程的角度ViT版本也比窗口注意力变体更容易读懂梯度因为每个token都直达全图backward路径没有窗口边界调试时少一层玄学。3. 环境搭建与数据预处理先让模型跑起来再说3.1 去雾数据集怎么组织训练、验证、测试目录结构与格式常见做法是准备三个目录train存放成对的模糊图和清晰图val用来每N个epoch算PSNRtest放真实雾图只用清晰图做参考或不做参考。为解决合成雾和真实雾的domain gap实际工程会把合成雾数据当作主训练集再用少量真实雾图做微调。目录组织示例data/ ├── train/ │ ├── hazy/ # 模糊图jpg 或 png │ └── clear/ # 对应的清晰图文件名与 hazy 一一对应 ├── val/ │ ├── hazy/ │ └── clear/ └── test/ └── real_world/ # 真实雾图无配对清晰图文件逐一配对是这里最容易翻车的地方。我见过不止一次因为train和clear里文件名顺序不一致模型学到的是「把上一张图变成下一张图的颜色」PSNR看着高换新图全糊。所以建议加载时用dict按文件名建立映射而不是拿两个列表zip。此外建议把所有图统一缩放长边不超过1024训练时用随机裁剪不要整图送入ViT——整图输入会让patch数量激增显存直接爆掉。验证集计算PSNR时要保持输入尺寸一致不能用训练时的随机裁剪否则每次验证的裁剪位置不同指标波动大。常见做法是固定中心裁剪到256×256或者对整图做滑窗推理再拼回来。我一般会把val集按浓度分三组——轻雾、中雾、重雾分别统计PSNR避免一组数据拉高平均值掩盖重雾区域的失效。3.2 预处理参数归一化、patch_size和数据增广的取舍ViT对输入有一个硬性要求图的高和宽必须是patch_size的整数倍。所以核心参数先定patch_size。去雾任务里我一般不用ViT-B/16那个16而是用patch_size8后面接一层浅卷积做下采样任务。这会让token数量变为16倍显存上升但对细节恢复明显更好。如果显存吃紧可以用16先跑通再降到8微调。归一化直接用视觉预训练里最通用的mean[0.485, 0.462, 0.406]、std[0.229, 0.224, 0.225]虽然这是分类任务的统计量但预训练权重就是按这个范围对齐的迁移效果最好自己重新算一套统计量反而会让预训练失效。注意模糊图和清晰图必须使用相同的mean/std。如果分别归一化相当于把清晰图的亮度基准改掉了模型永远学不到正确的颜色映射。数据增广方面随机水平翻转、随机裁剪到256×256或320×320、色彩抖动亮度/对比度/饱和度各±0.2就够。去雾任务不建议做随机擦除或强旋转因为雾的物理分布会被破坏模型会学歪。训练时把模糊图和清晰图用同一个随机种子做相同的裁剪与翻转保证像素严格对齐。这里我用一个自定义Dataset来实现配对增广而不是直接在transform里各做各的。3.3 训练入口的骨架最小可运行的启动脚本把训练循环封装成一个train.py入口部分大概是这个样子import torch from torch.utils.data import DataLoader from dataset import DehazeDataset from model import DehazeViT def main(): torch.manual_seed(0) device cuda if torch.cuda.is_available() else cpu train_ds DehazeDataset(data/train, patch_size8, crop_size256) train_dl DataLoader(train_ds, batch_size8, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastTrue) model DehazeViT(patch_size8, embed_dim192, depth12, num_heads3).to(device) optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) for epoch in range(1, 301): for hazy, clear in train_dl: hazy hazy.to(device) clear clear.to(device) pred model(hazy) loss criterion(pred, clear) optimizer.zero_grad() loss.backward() optimizer.step()这个骨架里crop_size统一了输入分辨率patch_size必须在数据集类和模型类里保持一致。batch_size8是ViT-B/8配置下比较平衡的值显存不够先减到4同时开梯度累积。优化器AdamW配合weight_decay1e-4是Transformer惯用组合和CNN常用的SGDmomentum差别很大直接套CNN的经验常常不收敛。学习率调度器我推荐cosine退火从1e-4降到1e-6总训练轮数300到400。ViT训练慢主要体现在前期前50个epoch可能看不出明显效果不要因为验证集没涨就急着降低学习率先确认loss在下降。训练日志里除了打印loss建议每50步打印一次当前学习率和显存占用方便定位训练停滞是因为学习率还是因为OOM。4. 模型实现细节从Patch Embedding到去雾输出头4.1 核心结构把ViT改造成回归模型的三个改动去雾ViT不能直接拿来分类要做三处改动。第一输入头用卷积做Patch Embedding而不是单纯切片重叠卷积能缓解patch边缘的不连续感。第二位置编码因为训练裁剪尺寸固定而长度固定推理不同尺寸时需要插值处理。第三输出头把Encoder输出reshape回特征图再接几个上采样卷积恢复原分辨率。简化实现如下import torch import torch.nn as nn class PatchEmbed(nn.Module): def __init__(self, in_ch3, embed_dim192, patch_size8): super().__init__() # 重叠卷积stride 取 patch_size 的一半保留相邻 patch 重叠信息 self.proj nn.Conv2d(in_ch, embed_dim, kernel_size3, stridepatch_size // 2, padding1) def forward(self, x): x self.proj(x) # (B, C, H/4, W/4) 当 patch_size8 时 x x.flatten(2).transpose(1, 2) return x # (B, L, C) class DehazeHead(nn.Module): def __init__(self, embed_dim192, out_ch3): super().__init__() self.head nn.Sequential( nn.Conv2d(embed_dim, embed_dim, 3, padding1), nn.ReLU(inplaceTrue), nn.Conv2d(embed_dim, out_ch, 3, padding1), ) def forward(self, x, h, w): B, L, C x.shape x x.transpose(1, 2).reshape(B, C, h, w) return self.head(x)Patch Embedding用卷积做的好处是保留局部纹理stride取patch_size一半时相邻patch有重叠输出特征图的高宽是输入的1/4而不是1/8给输出头保留了更多空间细节。输出头最简可以只接两层3×3卷积直接回归RGB值实测效果和带亚像素上采样的版本差不了太多但建议用双线性插值把特征图放大到原始尺寸再卷积边缘更干净。Transformer Encoder的depth和num_heads不用刻意调大。depth12、embed_dim192、num_heads3这个组合是显存和效果之间的平衡点。num_heads不是越多越好每个head的维度通常建议保持64左右192维拆3个head正好。qkv的bias建议保留对回归任务有一点稳定作用虽然分类任务里经常去掉。4.2 损失函数怎么配L1还是MSE感知损失加不加去雾训练的损失组合决定最终画质。纯MSE收敛快但图偏平滑细节糊纯L1能保留更多锐利边缘但训练早期梯度震荡。常见做法是L1为主加上SSIM损失和感知损失。感知损失不直接对VGG输出算L1而是取预训练分类网络几个中间层特征分别算L1再加权低层权重小、高层权重大这样既约束颜色又约束结构。def dehaze_loss(pred, target): l1 torch.mean(torch.abs(pred - target)) ssim_loss 1 - ssim(pred, target) perc perceptual_loss(pred, target) # 多层特征 L1 加权 return 1.0 * l1 0.5 * ssim_loss 0.1 * perc权重不是拍脑袋定的。先全部跑1个epoch看每个分量的量级把SSIM和感知损失的量级拉到和L1接近再微调系数。比如L1在0.1量级SSIM损失在0.05量级感知损失如果只有0.01那感知损失权重0.1就是合理的如果感知损失初始就有0.1权重需要降到0.01以下。如果发现训练后期PSNR不再涨优先检查是不是L1权重太大把感知损失压没了。4.3 训练循环与Checkpoint策略训练循环需要留意三点混合精度、EMA和checkpoint的保存策略。ViT训练很吃显存自动混合精度能省将近一半显存同时加速。EMA是对模型参数做滑动平均推理时用EMA参数比用step后的参数更稳去雾这类回归任务收益非常明显。Checkpoint不要只存model.state_dict()还要存optimizer、epoch、EMA参数、当前学习率否则恢复训练时优化器状态丢失学习率调度不平滑。scaler torch.cuda.amp.GradScaler() ema_model copy.deepcopy(model) ema_alpha 0.999 for epoch in range(1, 301): model.train() for hazy, clear in train_dl: hazy, clear hazy.to(device), clear.to(device) optimizer.zero_grad() with torch.cuda.amp.autocast(): pred model(hazy) loss dehaze_loss(pred, clear) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() with torch.no_grad(): for p, ema_p in zip(model.parameters(), ema_model.parameters()): ema_p.data.mul_(ema_alpha).add_(p.data, alpha1 - ema_alpha) if epoch % 50 0: torch.save({ epoch: epoch, model: model.state_dict(), ema: ema_model.state_dict(), optimizer: optimizer.state_dict(), scaler: scaler.state_dict(), }, fcheckpoints/dehaze_vit_{epoch}.pt)EMA固定alpha0.999在小数据集够用但严格做法是按step数修正alpha min(0.999, (step 1) / (step 10))训练初期让EMA跟得快一点。保存间隔建议前50个epoch每10个epoch存一次后面每50个存一次这样早期翻车还能找回能用的中间权重。恢复训练时要把scaler也恢复进去否则混合精度状态不一致偶尔会出现loss突然跳到NaN又恢复正常的情况看起来像玄学其实只是GradScaler的state丢了。5. 去雾训练与推理最常见的5个坑现象、原因与解决5.1 Loss停在0.1附近PSNR卡在22-23dB不再上涨现象训练到约50个epochL1 loss在0.1左右徘徊验证集PSNR长期停在22-23dB怎么加epoch都没用。原因多半是全局学习率不合适ViT对学习率比CNN敏感得多1e-4配合AdamW如果配错了调度器后期步长太大在局部震荡另外输入图像范围没有对齐如果图像是0-255而模型输出被sigmoid限制在0-1loss永远降不到位。排查方法很简单打印模型输出的像素值范围看是否和target一致。解决先把数据统一到0-1范围再用cosine退火调度器初始lr降到3e-5左右重新训练。跑两三个epoch看验证PSNR是否有明显起色再决定要不要继续。如果两三个epoch后PSNR还在22dB以下检查数据配对的文件名是否一致这个坑会让人白跑一周。5.2 浓雾区域出现色块和伪影现象恢复出来的图远处浓雾区域有规则色块边界和patch边界重合像马赛克一样。原因Patch Embedding直接切patchpatch之间没有信息交换位置编码每块一样导致patch边缘不连续输出头太浅没法缝合。这个现象在patch_size16时特别明显patch越大色块越粗。解决把Patch Embedding改成重叠卷积输出头加一层上采样卷积或者在Transformer Encoder后面接一个3×3的卷积修边。调完再看浓雾区域的PSNR一般能提升1-1.5dB。如果色块只出现在某些固定位置还要检查位置编码是否在训练时被冻结或者被错误广播。5.3 推理时换了图片尺寸直接崩溃现象训练时图是256×256推理时传入512×384报position embedding维度不匹配或者shape错误。原因ViT固定长度位置编码只支持训练尺寸。这是ViT系列最经典的坑任何用预训练ViT做下游任务的项目都会遇到。解决推理前把图padding到patch_size的整数倍再插值位置编码。常见做法是ph round(H / patch_size) * patch_size pw round(W / patch_size) * patch_size x F.interpolate(x, size(ph, pw), modebilinear, align_cornersFalse) # 推理完成后把输出裁剪回 (H, W)更彻底的做法是在模型forward里对位置编码做双线性插值这样任意尺寸都能跑代价是每换一次尺寸多一次小计算。我倾向于后者把逻辑内聚到模型里这样导出ONNX时也能一起固化。5.4 显存溢出但batch_size明明设得很小现象batch_size4、分辨率256×256还是OOM而同样配置的CNN完全没压力。原因ViT自注意力的复杂度是O(n²)ntoken数量。patch_size8时256×256有1024个token每个token要和1024个token计算注意力显存消耗远超CNN。此外混合精度没开也是常见助推因素一个batch的反向传播会把所有中间注意力矩阵都保存在显存里。解决优先开启AMP梯度缩放其次把patch_size从8改到16token数降为原来的1/4再不行就在Encoder层之间插入窗口注意力或者用梯度累积把batch_size降到1-2。梯度累积在实现上要留意loss除以累积步数后再backward否则等效学习率会变大模型训练不稳定。5.5 夜间和强光源场景全面失效现象白天有雾图效果很好换成夜间、路灯或车灯强光源场景输出偏色严重暗部死黑。原因大气散射模型假设大气光是全局均匀的夜间有局域强光源这个假设不成立模型在合成雾数据上学到的分布里没有这类样本属于分布外输入。解决如果只是白天场景落地直接限制使用范围如果非得做夜间要在模型里加一个物理先验分支把大气光A改成可学习的空间变化大气光图并在训练集里加入夜间强光源合成样本。这个改动对工程复杂度是数量级的提升我一般会先和需求方确认清楚夜间是否真有硬性要求。很多时候做演示用的夜间去雾最后用图像增强和色调映射就解决了不必硬上物理模型。6. 进阶把去雾ViT导出成ONNX做一套不依赖训练框架的推理脚本模型训练完要落地不能每次推理都开着Python训练环境。常见做法是导出ONNX。导出前先确认三点模型处于eval模式把位置编码插值逻辑固化到forward里避免导出时动态shape导致算子不支持输入输出都能用固定维度描述。用torch.onnx.export把模型导出为dehaze.onnxopset_version选11以上因为低版本对interpolate等算子的支持不稳定。import numpy as np import onnxruntime as ort from PIL import Image sess ort.InferenceSession(dehaze.onnx, providers[CUDAExecutionProvider, CPUExecutionProvider]) def dehaze_infer(img_path, patch_size8): img Image.open(img_path).convert(RGB) H, W img.size[1], img.size[0] ph round(H / patch_size) * patch_size pw round(W / patch_size) * patch_size img img.resize((pw, ph), Image.BILINEAR) arr np.asarray(img, dtypenp.float32) / 255.0 arr (arr - [0.485, 0.462, 0.406]) / [0.229, 0.224, 0.225] x arr.transpose(2, 0, 1)[None] # (1,3,H,W) out sess.run(None, {sess.get_inputs()[0].name: x})[0] out out[0].transpose(1, 2, 0) out out * [0.229, 0.224, 0.225] [0.485, 0.462, 0.406] out np.clip(out, 0, 1) * 255.0 return out[:H, :W].astype(np.uint8)这套脚本只依赖numpy、onnxruntime和PIL生产环境不用装PyTorch。我在实际中吃过亏的地方是归一化细节训练时用什么mean/std推理必须完全一致很多代码把这一步省了导致输出整体偏色。还有一处要注意ONNX Runtime的输入节点name在不同导出版本里会变用sess.get_inputs()[0].name而不是写死一个字符串能少踩一个坑。我现在的习惯是训练脚本和推理脚本共用一个preprocess.py所有归一化和padding逻辑都从同一处引用。验证模型是否导出正确拿同一张图分别跑PyTorch和ONNX输出对比最大像素差不超过1就说明转换无损。一套去雾工程真正可靠的验收方式不是看训练时PSNR而是拿一批现场真实雾图跑一遍让没参与建模的人看输出是否自然、有没有偏色这一步比任何指标都更能发现问题。希望帮到你。本文还有配套的精品资源点击获取