简介这份资源面向计算机视觉方向的学习者与研究者提供一套基于Vision Transformer的图像去雾算法完整实现适合具备一定深度学习基础、希望深入理解Transformer在底层视觉任务中应用的中高级开发者。压缩包共340个文件约156.38MB以204个Python源码为核心辅以yaml配置、csv实验记录、ipynb笔记、png与gif可视化结果及md说明文档覆盖模型定义、训练脚本、数据处理与结果展示等环节。项目内含预训练权重与数据集划分说明训练参数集中在option.py中可通过--pretrain_weights指定权重路径--train_ps控制输入补丁大小便于复现与二次调参。目前已有1445人学习下载读者可据此掌握ViT去雾的网络结构、训练流程与实验记录方式并借助可视化文件快速验证效果适合作为课程设计、科研复现或算法改进的参考工程。1. 从一张雾天高速监控图说起Vision Transformer 图像去雾到底在做什么去年冬天帮朋友处理一批高速卡口抓拍图雾霾天拍出来的车牌基本糊成一团传统暗通道先验去雾跑完车牌边缘是清楚了但天空区域直接出现色块断层远处车灯的光晕被当成雾浓度估计进去整张图偏色严重。后来换成基于 Vision Transformer 的去雾方案同一批图里天空过渡自然、车牌字符可辨这才意识到注意力机制在全局雾浓度建模上的优势确实不是玄学。这篇笔记就围绕「基于 Vision Transformer 的图像去雾算法研究与实现」这个方向把原理选型、Python 源码结构、训练参数、推理部署和踩坑记录一次讲透。它适合已经会用 PyTorch 跑通基础 CNN 去雾、想进一步理解 Transformer 在底层视觉任务里怎么落地的人也适合拿到一份 python 源码但不知道从哪改起、参数怎么调的从业者。核心要解决的问题只有一个给定一张有雾图怎么用 ViT 结构稳定地恢复出清晰图并且这套 python 源码能真正跑起来、改得动、复现得了。2. Vision Transformer 做去雾的原理与选型为什么不是简单把 CNN 换掉2.1 去雾任务的本质与 ViT 的切入点图像去雾的物理模型通常写成 I(x) J(x)·t(x) A·(1 - t(x))其中 I 是有雾图J 是待恢复的清晰图t 是透射率A 是大气光。传统方法靠先验估计 t 和 ACNN 方法靠数据驱动直接回归 J 或者回归 t、A 再合成。问题在于雾的分布是全局相关的一片天空区域的雾浓度往往和画面另一侧的光源、景深有关系卷积核的感受野再大也是局部堆叠出来的层数不够就抓不住这种长程依赖。Vision Transformer 的切入点就在这里。它把图像切成固定大小的 patch每个 patch 线性映射成 token然后通过多头自注意力让任意两个 patch 直接建立联系。对于去雾来说这意味着天空区域的 token 可以直接参考远处光源 token 的信息来判断雾浓度而不是靠一层层卷积慢慢传过去。常见做法是采用编码器-解码器结构编码器用 ViT 提全局特征解码器逐步上采样恢复分辨率中间用跳跃连接补细节。也有把 ViT 作为骨干、后面接轻量卷积头的混合结构兼顾全局建模和局部纹理。选型时要清楚一点纯 ViT 去雾在数据量不足时容易过拟合因为自注意力没有卷积那样的归纳偏置。所以实际项目里要么用在大规模数据上预训练过的 ViT 权重做初始化要么用混合结构要么在训练时加较强的数据增强。这是决定你这套 python 源码能不能训出效果的第一道分水岭。2.2 源码目录结构与模块职责拿到一份「python源码项目介绍使用说明」的压缩包先别急着跑 train.py。我一般会先花十分钟把目录结构摸清楚因为去雾项目的模块划分直接决定了你改哪一块。典型结构大致是这样dehaze_vit/ ├── configs/ # 训练与推理配置 │ └── default.yaml ├── datasets/ # 数据集加载与预处理 │ ├── __init__.py │ └── dehaze_dataset.py ├── models/ │ ├── __init__.py │ ├── vit_encoder.py # ViT 编码器 │ └── dehaze_net.py # 整体去雾网络 ├── utils/ │ ├── metrics.py # PSNR / SSIM │ └── image_io.py ├── train.py ├── inference.py └── requirements.txt看几个关键点models/dehaze_net.py 里 forward 的输入输出形状决定你能否换分辨率datasets/dehaze_dataset.py 里是否做了成对读取有雾图和有雾对应的清晰图必须文件名对齐configs/default.yaml 里 patch_size、embed_dim、depth 这几个参数决定模型大小和显存占用。如果源码里没有 configs 而是把参数硬编码在 train.py那改动成本会高一些建议先重构成配置文件再动手。2.3 环境搭建与最小可跑命令环境这块PyTorch 版本和 CUDA 匹配是老生常谈但每次都要确认的事。我一般用 conda 建独立环境避免和系统里的包打架conda create -n dehaze_vit python3.9 -y conda activate dehaze_vit # 根据你的 CUDA 版本装对应 torch这里以 cu118 为例 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install opencv-python pillow numpy pyyaml tqdm tensorboard装完先验证 GPU 可用import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else CPU only)如果 cuda.is_available() 返回 False别急着改代码先查驱动版本和 torch 编译版本是否匹配。这一步翻车的概率比模型本身还高。确认环境没问题后用一张图跑一次推理验证源码链路是通的python inference.py --input test_images/hazy_01.jpg --output results/ --ckpt weights/dehaze_vit.pth推理脚本一般会做三件事读图归一化、送进模型、反归一化保存。如果这一步报形状不匹配八成是推理时的输入尺寸和训练时的 patch 切分逻辑不一致后面避坑章节会细说。3. 训练自己的去雾模型数据、参数与损失函数怎么定3.1 数据集准备与成对样本对齐去雾是有监督任务训练集必须是有雾图和无雾清晰图成对出现。常用公开数据集有 RESIDE、Haze4K 这类但如果你有自己的业务图比如监控抓拍那自建数据集更贴合实际分布。自建时最容易出问题的是文件名对齐hazy 目录下 0001.jpg 必须对应 clear 目录下 0001.jpg一旦错位模型学出来的就是错误映射loss 可能还在降但推理结果一塌糊涂。我一般会写个校验脚本先过一遍import os hazy_dir data/train/hazy clear_dir data/train/clear hazy_files sorted(os.listdir(hazy_dir)) clear_files sorted(os.listdir(clear_dir)) assert len(hazy_files) len(clear_files), 数量不一致 mismatch [f for f, c in zip(hazy_files, clear_files) if f ! c] if mismatch: print(文件名不匹配示例:, mismatch[:5]) else: print(成对校验通过共, len(hazy_files), 对)这段逻辑很直白数量对齐、逐名对齐。参数上没什么可调的但建议把校验做成训练前的固定步骤而不是靠人眼抽查。数据量方面ViT 结构对数据量敏感少于几千对的话要么用预训练权重要么考虑混合 CNN 结构否则很难收敛到可用水平。3.2 关键训练参数patch_size、embed_dim、depth 与学习率ViT 去雾的效果和显存占用基本由这几个参数决定。下面这张表是我在 24G 显存上实测比较稳的一组配置供参考参数含义小数据集建议大数据集建议patch_size切块大小816embed_dimtoken 维度256512depth编码器层数612num_heads注意力头数88batch_size批大小48lr初始学习率1e-42e-4epochs训练轮数200100patch_size 越小token 数量越多全局建模越细但显存和计算量平方级上升。小数据集用 patch_size8 配合较浅的 depth能缓解过拟合。学习率我一般用余弦退火加 warmup前 5 个 epoch 线性升温避免一开始就把预训练权重冲坏。import torch from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) warmup LinearLR(optimizer, start_factor0.1, total_iters5) cosine CosineAnnealingLR(optimizer, T_max195, eta_min1e-6) scheduler SequentialLR(optimizer, schedulers[warmup, cosine], milestones[5])AdamW 的 weight_decay 别设太大1e-4 是比较稳的起点。如果训练 loss 震荡厉害先降学习率而不是加 batch_size这是血泪经验。3.3 损失函数组合L1、感知损失与 SSIM 的取舍只用 L1 损失训出来的去雾图PSNR 可能不低但视觉上偏平滑细节发糊。常见做法是 L1 为主加一点感知损失VGG 特征匹配和 SSIM 损失做约束import torch.nn as nn class DehazeLoss(nn.Module): def __init__(self, w_l11.0, w_ssim0.5, w_perc0.1): super().__init__() self.w_l1 w_l1 self.w_ssim w_ssim self.w_perc w_perc self.l1 nn.L1Loss() def forward(self, pred, target, perc_lossNone, ssim_lossNone): loss self.w_l1 * self.l1(pred, target) if ssim_loss is not None: loss loss self.w_ssim * ssim_loss if perc_loss is not None: loss loss self.w_perc * perc_loss return loss权重不是拍脑袋定的L1 保底SSIM 拉结构相似度感知损失补纹理。感知损失权重给太大会引入 VGG 自身的伪影0.1 左右比较安全。如果训练资源有限只用 L1 SSIM 也能出可用的结果感知损失属于锦上添花。4. 推理与部署把训练好的 ViT 去雾模型跑出稳定结果4.1 单图推理与批量推理的写法差异单图推理和批量推理看着差不多坑却不一样。单图通常直接 resize 到训练分辨率批量推理如果每张图尺寸不同要么统一 resize可能变形要么按原图尺寸做 padding 到 patch 的整数倍。我一般推荐后者避免变形影响去雾质量import torch import cv2 import numpy as np def pad_to_multiple(img, multiple16): h, w img.shape[:2] ph (multiple - h % multiple) % multiple pw (multiple - w % multiple) % multiple padded cv2.copyMakeBorder(img, 0, ph, 0, pw, cv2.BORDER_REFLECT) return padded, (h, w) def infer(model, img_path, devicecuda): img cv2.imread(img_path)[:, :, ::-1] / 255.0 padded, (h, w) pad_to_multiple(img, 16) tensor torch.from_numpy(padded).permute(2, 0, 1).unsqueeze(0).float().to(device) with torch.no_grad(): out model(tensor).clamp(0, 1) out out.squeeze(0).permute(1, 2, 0).cpu().numpy() out (out[:h, :w] * 255).astype(np.uint8) return out[:, :, ::-1]pad_to_multiple 的 multiple 要和模型下采样倍数一致ViT 加解码器通常是 16 或 32。裁回原尺寸这一步不能省否则输出图会比原图大一圈。4.2 显存不够时的三种降级策略推理时显存不够常见三种处理一是降分辨率把长边限制到 1024 再推理二是分块推理把大图切成重叠小块分别去雾再拼接三是半精度推理。分块推理要注意块间重叠否则拼接处会有明显接缝def tiled_infer(model, img, tile512, overlap64, devicecuda): h, w img.shape[:2] output np.zeros_like(img, dtypenp.float32) count np.zeros((h, w, 1), dtypenp.float32) step tile - overlap for y in range(0, h, step): for x in range(0, w, step): y2, x2 min(y tile, h), min(x tile, w) patch img[y:y2, x:x2] tensor torch.from_numpy(patch).permute(2, 0, 1).unsqueeze(0).float().to(device) with torch.no_grad(): out model(tensor).squeeze(0).permute(1, 2, 0).cpu().numpy() output[y:y2, x:x2] out count[y:y2, x:x2] 1 return output / np.maximum(count, 1)overlap 给 64 是经验值太小接缝明显太大推理变慢。半精度用 torch.cuda.amp.autocast 包一下即可但要注意有些 ViT 实现在 fp16 下 attention 会溢出出现 NaN 就退回 fp32。4.3 用 PSNR 和 SSIM 验证去雾效果训练日志里的 PSNR 只能参考真正判断效果要在固定测试集上跑一遍。PSNR 高不代表视觉好SSIM 更贴近结构保真度两个一起看from skimage.metrics import peak_signal_noise_ratio as psnr from skimage.metrics import structural_similarity as ssim def evaluate(pred, gt): p psnr(gt, pred, data_range255) s ssim(gt, pred, channel_axis2, data_range255) return p, s如果 PSNR 正常但 SSIM 偏低通常是结构信息丢失检查是不是感知损失权重过大或者解码器上采样方式太粗糙。反过来 SSIM 高 PSNR 低可能是整体偏色查一下归一化和反归一化是否对称。5. 避坑与排查ViT 去雾项目里最容易翻车的五件事5.1 训练 loss 下降但推理全黑或全白现象训练日志 loss 稳定下降推理输出却是纯黑或纯白。原因通常是归一化不一致——训练时把图归一化到 [0,1] 或 [-1,1]推理时忘了做同样处理或者反归一化时 clamp 范围写错。解决把训练和推理的预处理写成同一个函数两边共用别各写一份。5.2 显存溢出但 batch_size 已经调到 1现象batch_size1 仍然 OOM。原因多半是输入分辨率太高token 数量随分辨率平方增长。解决先降训练分辨率或者增大 patch_size 减少 token 数再不行就上梯度累积模拟大 batch。5.3 输出图有网格状伪影现象去雾结果出现规则网格纹路。原因是 patch 切分后 token 之间缺乏重叠解码器上采样时块边界没对齐。解决在 patch embedding 时用重叠切分或者解码器里加一层平滑卷积也可以在损失里加总变差约束抑制块效应。5.4 天空区域出现过曝或色偏现象天空本来该是均匀灰白结果出现亮斑或偏蓝偏黄。原因是大气光估计被高亮区域带偏ViT 的全局注意力把光源 token 的权重放太大。解决训练数据里增加天空占比高的样本损失函数里对高亮区域加权或者在网络输出后接一个颜色校正分支。5.5 换自己的数据集后效果骤降现象公开数据集上指标很好换成业务图就崩。原因是域差异——雾的浓度分布、成像设备、光照条件都变了。解决先用少量业务图做微调学习率调小到 1e-5冻结编码器前几层只训解码器和最后几层编码器。别一上来就全量重训容易把预训练学到的通用特征冲掉。6. 进阶技巧用注意力图定位雾浓度顺带验证模型是否真的学到了东西模型跑通之后我习惯做一件事把 ViT 最后一层的注意力图可视化出来看看它到底在关注哪里。这一步不是为了发论文而是为了判断模型是不是真的在学雾浓度分布还是靠数据集里的捷径特征蒙混过关。具体做法是注册一个 forward hook把某个注意力头的输出抓出来按 token 位置还原成热力图叠加到原图上import torch import matplotlib.pyplot as plt attn_maps [] def hook_fn(module, input, output): # 假设 output 形状为 [B, heads, N, N] attn_maps.append(output.detach().cpu()) # 假设模型里注意力模块叫 model.encoder.layers[-1].attn handle model.encoder.layers[-1].attn.register_forward_hook(hook_fn) _ model(sample_tensor) handle.remove() attn attn_maps[0][0] # 取第一个样本 attn attn.mean(dim0) # 多头平均 cls_attn attn[0, 1:] # 以 cls token 为查询去掉自身 grid int(cls_attn.shape[0] ** 0.5) heat cls_attn.reshape(grid, grid).numpy() plt.imshow(heat, cmapjet) plt.colorbar() plt.title(Attention heatmap) plt.show()看热力图时有几个判断标准如果高亮区域集中在天空和远景说明模型确实在关注雾浓度高的地方这是好现象如果高亮集中在画面边缘或者固定角落那大概率是数据集里存在位置偏差模型走了捷径换一批图就会失效。这个技巧帮我省过好几次返工——有一次发现模型只盯着画面右下角查了半天才发现训练集里右下角有个固定水印模型把水印当成了去雾线索。另一个进阶用法是把注意力图当作后处理权重。雾浓度高的区域注意力权重也高可以在融合多尺度输出时给这些区域更大权重细节区域则保留更多原始信息。这个思路不需要改网络结构只在推理后处理阶段加十几行代码对提升视觉观感有立竿见影的效果。最后说个习惯每次改完参数或者换数据集我都会固定跑三张图——一张天空占比大的、一张近景细节多的、一张夜景有光源的肉眼对比之后再去看指标。指标会骗人眼睛不会。这套流程坚持下来比盲目调参省的时间多得多。希望帮到你。本文还有配套的精品资源点击获取