简介基于Vision Transformer的图像去雾算法研究与实现源码与文档包专为计算机视觉方向的学生、科研人员及算法工程师设计围绕Uformer等Transformer结构在图像去雾任务中的应用展开。资源提供完整的Python工程包含NH-HAZE数据集预处理模块可将训练大图切分为256×256小图、训练主程序My_train.py、超参数配置nepoch270、batch_size32等及配套说明文档能够复现从数据准备、模型训练到结果分析的全流程。包体共338个文件其中204个py源码、16个yaml配置、12个csv训练指标、9个ipynb分析笔记、9个txt说明及png/gif可视化图等压缩包约156.36MB目录规划清晰便于按模块检索与二次修改。目前已有291人学习下载适合希望快速完成去雾实验、理解Vision Transformer在前向与反向传播中细节的进阶开发者。1. 基于Vision Transformer的图像去雾算法用ViT换掉U-Net的编码器值不值雾天拍出来的图对比度低、细节糊成一片。过去几年大家默认用CNN做去雾AOD-Net、FFA-Net这些模型在RESIDE数据集上刷榜落地也稳。但CNN的卷积核视野有限雾的浓度分布又往往是全局的——远处一片白茫茫近处还算清楚这种全局退化用局部卷积去拟合总有点吃力。Vision TransformerViT把图像切成patch后用全局注意力建模天然适合处理这种全局退化问题。所以基于Vision Transformer的图像去雾算法本质上就是把原来U-Net或残差网络里的卷积主干换成Transformer结构用自注意力去捕获雾的全局分布再回归出清晰图像。这类项目现在很常见打包成“源码文档说明”的Python工程适合正在做毕业设计、课程设计或者准备把去雾当切入点上手ViT的人。它解决的不是“能不能去雾”的问题而是“怎么用Transformer把去雾效果做得比CNN更好”的问题。你需要的是一台有NVIDIA显卡的电脑、PyTorch环境以及一份能跑通的代码。下面我按自己复现这类项目的经验把原理、训练、评估和踩坑一次性讲清楚。2. 从CNN到ViT去雾任务的原理与实际选型理由2.1 去雾的核心是求解大气散射模型不是“看图猜清晰图”图像去雾的物理基础是大气散射模型I(x) J(x)t(x) A(1 − t(x))。I是雾图J是清晰图t是透射率A是全局大气光。模型要做的事就是根据I反推J。传统方法用暗通道先验估计t和A再带入公式求解深度学习方法则是让网络直接回归J把物理模型的求解过程隐式学进网络参数里。当网络结构从CNN换成ViT时变化的核心在于感受野。CNN要看到远处像素必须靠堆叠卷积层或扩大膨胀率Transformer的注意力机制一步到位任何位置的像素之间都能直接建立关联。这对去雾的意义很大——雾带来的退化是跨区域的近处清晰区域的信息可以辅助恢复远处模糊区域这种“全局参考”能力正是CNN所欠缺的。2.2 为什么偏要用Vision Transformer三个具体的理由第一个理由是全局上下文建模。去雾质量差的区域往往纹理少、颜色偏灰白单靠局部邻域很难判断原来的颜色和纹理。ViT的全局注意力让每个patch在恢复时参考整张图的信息颜色偏移的纠正会更稳。第二个理由是透射率估计更平滑。透射率图在空间上是大片缓变区域CNN输出容易产生块状伪影而Transformer输出的特征图天然带有全局一致性透射率估计更平滑。第三个理由是迁移学习的红利。ViT在ImageNet上预训练的权重可以直接拿来初始化去雾模型的编码器效果比自己从头训练好得多。这是DehazeFormer等模型能跑出好分数的关键原因之一——它用了卷积patch嵌入和相对位置编码专门适配了图像恢复任务的局部先验需求。2.3 去雾ViT和标准ViT的差异不只是把分类头换成回归头直接用标准ViT做去雾会翻车。标准ViT的patch是16×16对于去雾这种像素级任务太粗糙而且ViT的Patch Embedding和位置编码是为分类设计的没有考虑图像的局部连续性。常见的做法是改成8×8甚至4×4的小patch或者在网络前半部分用卷积下采样、后半部分用Transformer做全局建模——这叫做卷积嵌入的混合结构。做去雾的ViT还面临一个分辨率问题分类任务常用224×224输入但图像去雾输入一般是256×256或更大。位置编码必须支持插值否则换输入尺寸就直接报错。我一般会在训练脚本里把位置编码模式设成interpolate并且只在尺寸变化不大时用。2.4 选型对比直接告诉你怎么选结构方案参数量去雾效果推理速度适用场景纯CNNAOD-Net小一般快嵌入式设备、实时视频CNNTransformer混合中等好中通用去雾、学术研究纯ViT大不稳定慢不建议日常使用我在实际项目中倾向于第三种变体——也就是DehazeFormer那样的设计用卷积提取底层特征中间几个stage用Transformer block做全局建模最后用卷积上采样恢复。结构和U-Net类似但编码器的核心块换成了Transformer。3. 落地实操数据集准备、损失函数与训练配置3.1 从零跑通训练的最小代码框架拿到一个去雾项目第一件事不是改模型而是先把训练脚本跑通。下面是简化后的最小训练框架核心是构建数据集、定义损失、设置优化器。# train_dehaze.py import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision import transforms from models import DehazeFormer # 项目自带模型定义 from datasets import HazyDataset # 项目自带数据集类 # 数据增强随机裁剪 水平翻转 颜色抖动 transform transforms.Compose([ transforms.RandomCrop(256), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.1, contrast0.1), transforms.ToTensor(), ]) train_set HazyDataset(rootdata/RESIDE/train, transformtransform) train_loader DataLoader(train_set, batch_size8, shuffleTrue, num_workers4) model DehazeFormer(embed_dim96, depths[1, 2, 8, 4]) model model.cuda() # 损失组合L1 感知损失权重各占一半 loss_l1 nn.L1Loss() loss_perceptual PerceptualLoss() # 项目里通常用VGG特征 optimizer torch.optim.AdamW(model.parameters(), lr2e-4, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max100, eta_min1e-6) for epoch in range(100): model.train() total_loss 0.0 for hazy, clean in train_loader: hazy, clean hazy.cuda(), clean.cuda() pred model(hazy) loss loss_l1(pred, clean) 0.5 * loss_perceptual(pred, clean) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm0.1) optimizer.step() total_loss loss.item() scheduler.step() torch.save(model.state_dict(), fcheckpoints/dehaze_epoch{epoch1}.pth) print(fEpoch {epoch1}: loss {total_loss / len(train_loader):.4f})这里的核心逻辑是每次迭代读一对雾图和清晰图模型输出预测的清晰图分别用L1损失和感知损失计算误差反向传播更新权重。clip_grad_norm_用来限制梯度范数防止ViT训练后期梯度爆炸——这是Transformer系模型的常见毛病。参数说明embed_dim96是第一层Transformer的通道数越小显存占用越少但表达力也会下降depths控制每个stage堆叠的Transformer block数量[1,2,8,4]是一个性价比不错的配置。学习率2e-4是AdamW在ViT系模型上的常见起点太大容易发散太小收敛极慢。3.2 数据集的坑RESIDE怎么划分训练集怎么选去雾领域的事实标准数据集是RESIDE但要注意它分好几个子集ITS用于合成雾训练、OTS用于户外场景、SOTS是测试集。做研究一般用ITS的13个场景训练在SOTS上评估如果数据量不够可以合并OTS。但真实雾图是没有配对清晰参考的只能靠合成雾图来训练。训练时一个容易被忽略的细节是合成雾的数据增强参数要在加载时实时生成不要提前把雾图存好。因为固定存储的雾图会让网络过拟合特定的透射率分布训练时动态随机生成雾度参数随机透射率t、随机大气光A能让模型的泛化能力明显上升。我在数据类里会这样处理# datasets.py 中的关键片段 class HazyDataset(Dataset): def __init__(self, root, transformNone): self.clean_paths glob.glob(root /*.png) self.transform transform def __getitem__(self, idx): clean cv2.imread(self.clean_paths[idx]) clean cv2.cvtColor(clean, cv2.COLOR_BGR2RGB) clean self._random_crop(clean, size256) # 动态合成雾图随机透射率和大气光 t np.random.uniform(0.6, 1.0, sizeclean.shape[:2])[..., None] A np.random.uniform(0.8, 1.0, size(1, 1, 3)) hazy clean * t A * 255 * (1 - t) hazy np.clip(hazy, 0, 255).astype(np.uint8) if self.transform: hazy self.transform(hazy) clean self.transform(clean) return hazy, clean这个数据类的好处是不需要预先生成雾图文件每次读取清晰图后随机生成透射率和大气光相当于无限数据增强。注意t的范围取0.6到1.0太小的透射率意味着浓雾过拟合浓雾会让模型在轻雾场景下过度增强。3.3 损失函数怎么搭配别迷信单一L2很多初学者直接上MSELoss训练完发现图片变模糊了——这是L2损失的通病像素级的均方误差会把所有高频细节磨平。去雾任务的常用搭配是L1损失打底加上VGG感知损失提升视觉质量再配合SSIM损失补充结构信息。感知损失不是自己去实现VGG而是用预训练好的VGG19提取中间层特征计算特征图的L1距离。这样网络学习的目标从“像素值和原图一样”变成“特征表达和原图一样”细节恢复效果明显更好。SSIM损失则强调局部结构一致性对边缘保持有好处。三者比例的常见起点是1:0.5:0.1但建议先用前两个训练等loss降下去再加SSIM损失。class CombinedLoss(nn.Module): def __init__(self, weights(1.0, 0.5, 0.1)): super().__init__() self.w_l1, self.w_perc, self.w_ssim weights self.l1 nn.L1Loss() self.perc PerceptualLoss() self.ssim SSIMLoss() def forward(self, pred, target): l1_loss self.l1(pred, target) perc_loss self.perc(pred, target) ssim_loss self.ssim(pred, target) return self.w_l1 * l1_loss self.w_perc * perc_loss self.w_ssim * ssim_loss3.4 优化器、学习率与Batch Size的设置建议Vision Transformer对超参数比CNN更敏感这是Transformer家族的共性。优化器建议用AdamW而不是SGD——ViT训练用SGD基本走不动。学习率方面一段式训练用CosineAnnealing从2e-4降到1e-6比较省心如果先用小学习率热身50个iteration再升到正常值训练会更稳这里直接在一个Epoch内完成热身即可。Batch Size按显存来12GB显存跑256×256输入batch_size8是极限如果OOM就把batch_size调到4同时把torch.backends.cudnn.benchmark设为True加速卷积。注意Batch Size减小后要适当降低学习率经验公式是学习率按batch size的平方根缩放——从8降到4学习率大约除以1.4。这个细节不处理小batch训练容易不收敛。3.5 训练资源的合理预期ViT系去雾模型比CNN慢很多。以我的经验一张12GB的RTX 3060输入256×256、batch size为4训练100个epoch大概需要3到5天具体要看模型深度和数据集规模。如果时间不够有两个选择一是加载ImageNet预训练权重收敛速度能快一倍以上二是把训练分辨率降到224×224最后再用256微调十几个epoch。4. 推理与量化评估PSNR/SSIM到底怎么算才算数4.1 用训练好的权重对单张图去雾训练跑完后需要把模型导出并进行推理。下面的代码展示了对一张雾图执行去雾推理并保存结果的完整流程。# inference.py import torch import cv2 import numpy as np from models import DehazeFormer device torch.device(cuda if torch.cuda.is_available() else cpu) model DehazeFormer(embed_dim96, depths[1, 2, 8, 4]) model.load_state_dict(torch.load(checkpoints/dehaze_epoch100.pth, map_locationdevice)) model.eval().to(device) def dehaze_image(model, img_path, output_path): img cv2.imread(img_path) origin_h, origin_w img.shape[:2] # 为了保证位置编码不失效先把输入缩放到训练尺寸的整数倍 scale 256 / min(origin_h, origin_w) new_h, new_w int(origin_h * scale), int(origin_w * scale) new_h, new_w new_h // 8 * 8, new_w // 8 * 8 # 8的倍数兼容patch划分 img_resized cv2.resize(img, (new_w, new_h)) img_tensor torch.from_numpy(img_resized.transpose(2, 0, 1)).float().div(255.0) img_tensor img_tensor.unsqueeze(0).to(device) with torch.no_grad(): dehazed model(img_tensor) dehazed dehazed.squeeze(0).cpu().numpy().transpose(1, 2, 0) dehazed np.clip(dehazed * 255, 0, 255).astype(np.uint8) dehazed cv2.resize(dehazed, (origin_w, origin_h)) cv2.imwrite(output_path, dehazed) if __name__ __main__: dehaze_image(model, test/hazy.jpg, test/result.jpg)这段代码做了三件关键事把输入图缩放到训练尺寸的整数倍以满足patch划分要求、前向推理后裁剪到0到255范围、最后缩回原图大小保存。整个过程不需要GT图任何一张雾图都能跑。4.2 PSNR和SSIM的计算方式不要直接调sklearnPSNR和SSIM需要在有配对清晰图的测试集上计算。PSNR公式里有个常见的细节像素值的取值范围MaxI取255还是1计算结果差异很大。很多项目为了刷分把MaxI取1但图像像素是0到255的uint8正确做法是MaxI255这叫PSNR不能跨实现直接比较的原因——同一张结果图不同库算出来能差好几个dB。SSIM的窗口大小也影响结果。默认的11×11高斯窗口是通用选择但去雾任务里如果要和论文对标得看作者用的是什么滤波窗口。通常我在报告里会写明用的scikit-image的structural_similarity并标注win_size11。评估脚本代码相对简单但重算一遍才能真正理解数据。# evaluate.py from skimage.metrics import peak_signal_noise_ratio, structural_similarity def evaluate(gt_dir, pred_dir): psnr_sum, ssim_sum, count 0.0, 0.0, 0 for name in os.listdir(pred_dir): gt cv2.imread(os.path.join(gt_dir, name)).astype(np.float64) pred cv2.imread(os.path.join(pred_dir, name)).astype(np.float64) psnr peak_signal_noise_ratio(gt, pred, data_range255) ssim structural_similarity(gt, pred, channel_axis-1, data_range255) psnr_sum psnr ssim_sum ssim count 1 print(fAverage PSNR: {psnr_sum / count:.2f} dB, SSIM: {ssim_sum / count:.4f})4.3 量化指标的陷阱PSNR高不代表视觉好这是去雾领域最典型的翻车现场PSNR刷到30以上但看图发现颜色偏灰、细节发糊。原因在于PSNR对整体像素误差敏感对局部结构不敏感——把整张图稍微调亮或调暗PSNR就会显著变化但视觉质量并没有变好。所以评估时必须两个指标一起看并配上视觉对比图。我的习惯是跑完测试后把三张图拼在一起看雾图原图、模型输出、GT清晰图。肉眼观察三个关键区域——天空区域有没有过曝、远处建筑物边缘是否清晰、近处植被纹理是否自然。如果PSNR高但天空变成了灰白色说明大气光估计有问题模型把天空区域错误增强了。5. 训练与部署避坑5个值得写进笔记里的踩坑记录5.1 训练loss一直不降在0.78附近震荡现象训练了十几个epochloss基本不动就像被卡住了一样。原因最常见的是学习率设置过大或过小。ViT对学习率很敏感AdamW在2e-4时正常调到5e-4就可能发散降到1e-5则几乎不动。还有一个输出层没有加Tanh或Sigmoid限制输出范围导致模型输出像素值溢出到负数或大于1L1 Loss在边界附近形成梯度震荡。解决把学习率拉回2e-4区间模型输出层确认是否有范围限制。如果是自己的实现在最后加一层nn.Tanh()把输出压缩到−1到1之间配合输入归一化到−1到1范围使用。5.2 显存OOMbatch size2都跑不动现象一开训练就报CUDA out of memory。原因ViT的注意力计算量随patch数量呈平方增长256×256输入切成8×8 patch序列长度是1024注意力矩阵就是1024×1024。显存占用大头是注意力矩阵和中间激活值不是参数本身。解决第一个方案是减小patch大小而不是降低分辨率保持分割块数不变会降低计算量把embed_dim从96降到64。第二个是开梯度检查点gradient checkpointing牺牲少量训练时间换取显存——PyTorch里用torch.utils.checkpoint.checkpoint包住Transformer block即可。第三个是把batch size降到1配合梯度累积模拟batch size为4的效果。5.3 PSNR高但视觉结果偏灰雾已经去了但颜色不对现象量化指标不错但输出图整体蒙了一层灰色天空区域尤其明显。原因训练数据里合成雾的全局大气光A取值偏向白色或亮灰而真实雾场景的大气光往往带点蓝色或黄色色偏。模型学到了“把像素往灰色推”的倾向就去雾的同时把色彩饱和度也削弱了。解决数据增强中加入大气光颜色的随机扰动不要让A固定在纯白色。把大气光取值范围设为[0.7, 1.0]并允许每个通道独立随机取值让模型见过多种色偏的雾。训练后再看结果色彩的灰蒙感通常会减轻。5.4 在真实雾图上效果差训练集PSNR很高现象合成雾数据集上PSNR有30一拿到真实拍摄的雾图去雾效果明显退化远处物体还是糊的。原因合成雾图和真实雾图的物理过程不完全一致。合成雾假设透射率均匀分布真实雾图的透射率和景深强相关近处透射率高、远处透射率低且存在多次散射和大气湍流等复杂因素。这就是领域差距问题。解决加入真实雾图进行微调。真实雾图没有GT可以用无参考指标做筛选后把网络输出周边一致性作为伪标签来微调。更简单的做法是用预训练模型在少量真实雾图上做风格迁移但效果有限。合理的预期是在目标域数据量有限时先用合成雾数据训练做主模型到了现场再针对性微调。5.5 推理速度太慢视频流实时处理跟不上现象单张256×256图片推理要200ms做不到实时处理。原因ViT的全局自注意力计算量太大多尺度结构叠加深了参数量。解决从轻量化入手。第一个方案是把分辨率降下来用160×160输入但质量会损失。第二个是换掉全局注意力用窗口注意力或局部注意力减少计算量这是Swin Transformer的思路。第三个方案是把训练好的ViT蒸馏到一个小型CNN里——教师模型是ViT学生模型用轻量U-Net蒸馏后的CNN推理速度能快5到10倍效果仍然比直接用CNN好。如果只是验证可行性项目演示时用GPU加速推理并缓存结果就够了别要求实时。6. 进阶技巧视觉结果自检与服务化部署量化指标算完不算完落地时还有一个容易忽略的环节视觉自检。我的习惯是每一轮训练结束挑3张有代表性的雾图出来跑推理一张看天空、一张看中景建筑、一张看近景植被。如果近景植被纹理清晰但天空变成灰白说明大气光估计引入偏置如果远景恢复出来了但近景出现振铃说明网络学的是锐化而不是去雾。这类主观检查花不了几分钟但能在PSNR没变的情况下看出模型的真实表现。另一个技巧是无参考条件下快速验证去雾效果对比去雾前后的暗通道均值。去雾算法的目标是让暗通道值恢复到0.1以下如果输出图的暗通道均值仍然很高说明雾没有去干净。这个脚本写起来很简单用opencv取每个像素的RGB最小值做最小值滤波再算均值即可。虽然没有GT图这么严谨但在真实场景里评估模型有没有正常工作很实用。如果项目要交付成服务可以封装成FastAPI接口输入图片URL返回去雾后的图片bytes也可以转成ONNX格式大幅提升加载速度和推理效率。转换时注意把位置编码设成固定buffer避免动态图和静态图的转换报错。这个技巧值得记录。作为一个反复复现过这类项目的人我的教训是不要迷信源码里的默认参数也不要一上来就换大模型。先把原始代码用最小数据集跑通确认loss下降、推理有输出再去调模型结构、补损失函数、上数据增强。ViT去雾的潜力在于全局建模能力局限在于计算量和数据需求理解这两点项目推进就不会走偏。希望这篇文章帮你少走几步弯路。本文还有配套的精品资源点击获取
