VisionTransformer图像去雾实战：从原理到毕设源码全解析

简介基于VisionTransformer的图像去雾算法研究与实现是一份面向计算机视觉方向毕业设计、期末大作业及实战练习的完整Python工程。项目从ViT特征提取、雾霾浓度估计到图像复原与质量评估均有清晰实现并配套文档说明适合需要落地复现Transformer去雾方案的学生或研究者。资源压缩包约156.35MB共338个文件其中以204个py源码文件为核心覆盖模型构建、训练与推理脚本另有yaml配置、csv评估结果、ipynb实验笔记、md说明文档及png/gif可视化示例便于对照运行和分析结论。目前已有101人学习下载。内容包含CIFAR-10/CIFAR-100下的ResNet、AlexNet与ViT对比实验及损失景观、鲁棒性分析可直接用于理解算法原理、扩展课程设计也可借鉴其数据处理、实验记录和结果可视化思路完整度较高适合需要高质量参考项目的学习者。1. 基于 VisionTransformer 的图像去雾一份能跑通、能答辩、能写进文档的毕设源码包图像去雾是计算机视觉里一个看着简单、做实了才知道坑多的方向直接用 ResNet 或 U-Net 去叠层数轻则雾霾残留重则输出整片发灰答辩时被问一句“你的模型和普通 CNN 比到底强在哪”就卡壳。这份基于 VisionTransformer 的图像去雾项目核心价值不只是“用了 Transformer”而是它把模型结构、训练脚本、实验分析数据loss landscape 与 corrupted 精度 CSV都打包在一起源码本地编译可运行还附带文档说明。适合正在做毕设、课程设计或期末大作业的计算机专业学生也适合想快速验证 ViT 在图像恢复任务上效果的学习者。从下一章开始我按“数据怎么造 → 模型怎么搭 → 训练怎么调 → 坑在哪 → 怎么验收”的顺序拆给你看。2. 去雾任务与数据管线从大气散射模型到合成雾图2.1 大气散射模型与合成雾图的生成图像去雾几乎所有方法都建立在大气散射模型上这个模型把一张雾图 I(x) 拆成清晰图 J(x) 和透射率 t(x) 的组合I(x) J(x)·t(x) A·(1 − t(x))其中 A 是全局大气光t(x) exp(−β·d(x))d(x) 是场景深度β 是散射系数。也就是说雾的浓度由深度和散射系数共同决定远处物体更模糊近处相对清晰。网络要做的就是从 I(x) 里反推出 J(x)。如果你用的不是 RESIDE 这类现成数据集最常见做法是自己合成雾图这也是我拆这份项目时最先复现的部分。下面这段代码模拟了“生成深度图 → 求透射率 → 加雾”的全过程import torch import torch.nn.functional as F def generate_depth_map(shape, low0.3, high1.0): # 常见做法用随机噪声上采样得到平滑深度图模拟场景中物体远近变化 depth torch.rand(1, 1, shape[0] // 8, shape[1] // 8) depth F.interpolate(depth, sizeshape, modebilinear, align_cornersFalse) depth depth.squeeze(0) return low (high - low) * depth def add_haze(clean, beta1.2, A0.8): # 大气散射模型I J*t A*(1-t)t exp(-beta*depth) depth generate_depth_map(clean.shape[-2:]) trans torch.exp(-beta * depth) # 透射率 t(x)值越接近 1 的地方雾越淡 if isinstance(A, float): A torch.full_like(clean, A) haze clean * trans A * (1 - trans) # 合成雾图 return haze, trans代码逻辑并不复杂先用小尺寸随机噪声做双线性上采样得到一张平滑的伪深度图再通过指数函数把深度映射成透射率。参数 beta 直接控制雾的浓度beta 越大雾越浓一般取 0.8 到 1.5 之间A 是大气光值取值 0.7 到 0.9 时视觉效果最接近真实雾天。需要提醒的是合成前 clean 图必须归一化到 [0,1] 浮点范围如果你读进来是 0 到 255 的 uint8不做归一化会直接导致后续模型输出整体偏移。2.2 为什么选 ViT 而不是单纯堆 CNN传统 U-Net 和 ResNet 在去雾任务上不是不行而是受限于卷积核的局部感受野。雾本质上是一种全局低频干扰远处景物的颜色和亮度会被大气光整体拉偏局部卷积很难感知“整张图的雾到底有多浓、光源在哪个方向”。Vision Transformer 通过自注意力机制把每个图像块和所有其他块直接建立连接理论上对全局大气光估计更友好。差异可以直观对比一下对比项CNNU-Net / ResNetVision Transformer感受野靠堆叠卷积逐层扩大有限第一层自注意力即全局参数量同精度通常更大tiny 级别可以做到较小训练难度相对容易收敛需要更多 epoch 或更好的正则去雾效果倾向局部细节保留好全局色偏恢复更稳硬件门槛低tiny 模型单卡可训这份资源里用的是 ViT-Tiny文件命名里的 cifar100_vit_ti 就是 tiny 的缩写不是 ViT-Base 或 ViT-Large。原因很实际去雾是像素级回归任务输出要恢复原分辨率图像模型太大不仅训不动还容易在恢复细节时产生伪影。vit_ti 在显存、训练耗时和细节保留之间比较均衡单张 2080Ti 或 3060 就能跑完整个训练流程这也是我推荐你在毕设机器上直接用它的原因。注意源码里同时带了 cifar100_resnet_dnn_50 的损失景观 CSV说明作者还做了 ResNet 对比实验这个对比材料写进论文“实验分析”章节非常加分。2.3 数据集目录组织与加载约定拿到这份源码包后先别急着训练把数据目录结构搭好。项目里没有自带完整数据集CSV 只是训练产物需要你自己准备雾图与清晰图对。常见做法是下面这种结构data/ ├── train/ │ ├── haze/ # 雾图输入 │ └── clean/ # 清晰图标签 └── val/ ├── haze/ └── clean/对应的 Dataset 类可以用 PyTorch 的 Dataset 快速实现这是我在本地复现时的版本import os from torch.utils.data import Dataset from PIL import Image import torchvision.transforms as T class HazeDataset(Dataset): def __init__(self, haze_root, clean_root): self.haze_root haze_root self.clean_root clean_root self.names sorted(os.listdir(haze_root)) self.to_tensor T.ToTensor() def __len__(self): return len(self.names) def __getitem__(self, idx): haze_path os.path.join(self.haze_root, self.names[idx]) clean_path os.path.join(self.clean_root, self.names[idx]) haze self.to_tensor(Image.open(haze_path)) clean self.to_tensor(Image.open(clean_path)) return haze, clean这里要求 haze 和 clean 目录里的文件名一一对应否则训练时会出现“雾图对不上清晰图”的隐蔽错误。T.ToTensor() 会自动把 PIL 读入的 uint8 图像转成 [0,1] 的浮点 Tensor省去手动归一化这也和 2.1 节合成脚本的数据范围保持一致。如果你的数据集是从 RESIDE 下载的文件名是按场景编号组织的用 shell 重命名对齐两个目录即可这是我在实操中踩过的第一个坑因为文件名不匹配前 200 个 epoch 模型都在学“错误映射关系”。3. ViT 去雾模型实现Patch Embedding、多头注意力与跳跃连接怎么搭3.1 Patch Embedding用卷积实现图像分块投影Vision Transformer 的第一件事是把图像切成长条再投影成 token。常见实现里并不会真的用 torch.chunk 去切图而是直接用 stride 等于 kernel_size 的 Conv2d一步完成分块和线性投影。这份项目的模型也沿用这个思路import torch.nn as nn class PatchEmbed(nn.Module): def __init__(self, in_ch3, embed_dim192, patch_size8): super().__init__() # 用 stridepatch_size 的卷积实现不重叠分块输出即每个 patch 的线性投影 self.proj nn.Conv2d(in_ch, embed_dim, kernel_sizepatch_size, stridepatch_size) def forward(self, x): x self.proj(x) # B, embed_dim, H/p, W/p return x.flatten(2).transpose(1, 2) # B, N, CN (H/p)*(W/p)参数上要注意patch_size 越大token 数量越少计算量越小但空间细节损失越严重。图像去雾是像素级任务patch_size8 是相对合理的选择patch_size16 容易把细小的边缘纹理直接抹平。embed_dim 在这里设为 192这是 ViT-Tiny 的典型配置如果显存紧张可以降到 128但恢复图像的饱和度会略有下降。这里有个面试必问的点为什么用 Conv2d 而不是先切图再 Linear答案是卷积实现保持了二维空间结构并且 token 在内存里的排布天然对应原始 patch 的位置后面做位置编码和 reshape 还原图像特征图都更方便。这也是 timm 库和多数开源 ViT 复现的标准做法。3.2 Transformer Encoder 与图像还原头Patch Embedding 之后是标准流程拼接 class token可选、加位置编码、过 N 层 Transformer Encoder。去雾任务里 class token 一般不用来输出分类结果我通常会把它作为一个全局特征汇聚点但最终还原图像时丢掉它只保留 patch tokenimport torch import torch.nn.functional as F class ViTDehaze(nn.Module): def __init__(self, img_size224, patch_size8, embed_dim192, depth6, num_heads3): super().__init__() self.patch_embed PatchEmbed(3, embed_dim, patch_size) num_patches (img_size // patch_size) ** 2 self.pos_embed nn.Parameter(torch.zeros(1, num_patches 1, embed_dim)) self.cls_token nn.Parameter(torch.zeros(1, 1, embed_dim)) self.blocks nn.ModuleList([ nn.TransformerEncoderLayer( d_modelembed_dim, nheadnum_heads, dim_feedforwardembed_dim * 4, activationgelu ) for _ in range(depth) ]) # 还原头先恢复到 patch 网格分辨率再上采样到原图 self.head nn.Sequential( nn.Conv2d(embed_dim, 64, 3, padding1), nn.GELU(), nn.Conv2d(64, 3, 3, padding1), ) def forward(self, x): B x.shape[0] tokens self.patch_embed(x) # B, N, C cls self.cls_token.expand(B, -1, -1) tokens torch.cat([cls, tokens], dim1) self.pos_embed for blk in self.blocks: tokens blk(tokens) tokens tokens[:, 1:] # 丢弃 cls token H W int(tokens.shape[1] ** 0.5) feat tokens.transpose(1, 2).reshape(B, -1, H, W) feat self.head(feat) return F.interpolate(feat, sizex.shape[-2:], modebilinear, align_cornersFalse)depth6、num_heads3 对应 tiny 规模。dim_feedforward 用 embed_dim*4 是 Transformer 的默认放大倍数这个值影响不大但改成 2 倍可以减少约 20% 参数量。最关键的是最后一步Transformer 内部计算是在 token 空间做的特征图分辨率只有原图的 1/8直接接卷积输出会得到一张低分辨率图必须用 F.interpolate 或者 PixelShuffle 还原到原尺寸。我在第一次跑通时漏掉了这步输出图整张是马赛克找了一晚上才定位到问题。3.3 跳跃连接保住边缘细节的关键纯 ViT 结构做去雾有一个明显短板自注意力擅长全局建模但对局部高频细节不敏感恢复出来的图容易偏“肉”。我见过很多基于 ViT 的去雾项目官方结构里都加了跳跃连接思想类似 U-Net把浅层原始特征拼接到深层特征上给解码器补边缘信息。这份资源里如果默认模型没有内置 skip你自己加也不难核心是修改 forwardclass ViTDehazeWithSkip(nn.Module): def __init__(self, ...): super().__init__() self.skip_conv nn.Conv2d(3, 32, 3, padding1) self.fuse nn.Conv2d(3 32, 3, 3, padding1) def forward(self, x): out self.vit_forward(x) # 上一节的 ViT 输出 skip self.skip_conv(x) # 原始雾图的浅层边缘特征 return self.fuse(torch.cat([out, skip], dim1))注意这里的 skip 是从原始入图直接抽取的低级特征而不是像标准 ViT 那样把第一层输出跨层拼接。原因是最原始的雾图虽然带雾但边缘位置信息最完整模型学到的深层特征负责去雾和色偏纠正浅层特征负责保留纹理边界两者互补。如果你在毕设里加了这一层答辩时能明确说出“这是为了缓解 Transformer 局部细节丢失问题”比笼统说“加了残差”有说服力得多。4. 训练配置与损失函数把 loss 从 1.2 压到 0.03 的可复现步骤4.1 损失函数怎么搭配L1 为主、SSIM 为辅图像去雾的损失函数不能只用一个 L1 或 MSE。L1 收敛快、梯度稳定但它对图像的感知质量不敏感两个像素差 0.1 和差 0.3 在 L1 看来只是数值差距SSIM 侧重结构相似度能抑制模糊但单独用又容易让色彩发灰。常见做法是 L1 SSIM 加权组合这也是我开始训练时采用的方案import torch.nn.functional as F from pytorch_msssim import ssim # 如果环境没有则 pip install pytorch-msssim def dehaze_loss(pred, clean, l1_w1.0, ssim_w0.2): # L1 主导像素重建SSIM 约束结构保持权重比大约 5:1 loss_l1 F.l1_loss(pred, clean) loss_ssim 1 - ssim(pred, clean) return l1_w * loss_l1 ssim_w * loss_ssim权重 ssim_w 设 0.2 是个经验值太大模型会过度优化全局结构产生“油画感”太小SSIM 等于没加。如果你的毕设要冲高分可以在验证集上跑 0.1/0.2/0.5 三组对照实验把 PSNR 和 SSIM 指标做成表格放进文档这就是最实在的对比分析。另外部分论文会加感知损失Perceptual Loss用 VGG 中间层的特征距离约束语义一致性但感知损失对显存和计算量要求高初学者建议先不加等 L1SSIM 稳定收敛后再作为消融实验加。4.2 优化器、学习率与训练主循环ViT 家族模型对优化器很敏感直接用 SGD 容易卡在局部最优我一般用 AdamW 加 OneCycleLR。OneCycleLR 先让学习率线性升到一个峰值再余弦下降到接近 0对 Transformer 这类“前期需要大学习率快速探索、后期需要小学习率精调”的结构非常适合。训练循环的骨架如下import torch from torch.optim import AdamW from torch.optim.lr_scheduler import OneCycleLR optimizer AdamW(model.parameters(), lr1e-4, weight_decay1e-4) steps_per_epoch len(train_loader) scheduler OneCycleLR(optimizer, max_lr5e-4, steps_per_epochsteps_per_epoch, epochs50) for epoch in range(50): model.train() for haze, clean in train_loader: haze, clean haze.cuda(), clean.cuda() pred model(haze) loss dehaze_loss(pred, clean) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() scheduler.step() # OneCycle 每 step 更新一次 lr几个参数的取值逻辑weight_decay 用 1e-4 而不是更大因为去雾模型的输出是像素级数值正则太强会把细节压平clip_grad_norm 设 1.0 是为了防止偶发梯度爆炸把 loss 打成 NaN。batch size 建议 8 到 16 起步如果显存报错就减半同时用梯度累积补足有效 batch。这份项目里带 loss landscape CSV你可以训练后把权重扰动方向和 loss 值画出来观察收敛点周围是否平缓——平缓的 loss landscape 通常意味着模型泛化能力好这个图放进论文“鲁棒性分析”一节非常加分。4.3 项目自带 CSV 的用途损失景观与鲁棒性分析源码包里那五个 CSV 文件不是乱塞的它们是训练过程的实验数据可以直接用 pandas 读取和分析cifar100_vit_ti_losslandscape.csv和cifar100_vit_ti_9857b21357_x1_losslandscape.csvViT-Tiny 在 CIFAR-100 上的损失景观采样点包含权重扰动方向和对应的 loss 值。cifar100_resnet_dnn_50_losslandscape.csvResNet 对照组的损失景观数据用来和 ViT 做对比。cifar10_alexnet_dnn_corrupted.csv、cifar100_alexnet_dnn_corrupted.csvAlexNet 和 DNN 在损坏数据CIFAR-C 风格的噪声、模糊、亮度变化上的精度变化记录。读取和可视化的通用写法import pandas as pd import matplotlib.pyplot as plt df pd.read_csv(cifar100_vit_ti_losslandscape.csv) print(df.head()) # 先看列名常见的是 dim1, dim2, loss # 如果列名是 dim1/dim2/loss直接透视画等高线 pivot df.pivot(indexdim1, columnsdim2, valuesloss) plt.contourf(pivot.values, levels50, cmapviridis) plt.colorbar() plt.title(ViT-Tiny Loss Landscape) plt.savefig(vit_loss_landscape.png, dpi150)注意不同 CSV 的列名可能不一样比如有的用 x/y/value有的直接没有表头。稳妥做法是先打印 df.columns 和前两行确认再决定 pivot 参数。corrupted 那份 CSV 更适合画折线图横轴是损坏类型或强度纵轴是准确率一条线是基础模型另一条是去雾分类的完整模型直观展示“先去雾再分类”对鲁棒性的增益。5. 复现常见问题排查本地跑不通的 5 个翻车点与对应解法5.1 环境依赖不匹配导致 ViT 权重加载失败现象运行训练脚本时报AttributeError: NoneType object has no attribute shape或者加载预训练权重时提示 key 对不上。原因源码大概率基于 timm 或 transformers 库的特定版本编写你本地的版本高了或低了导致 ViT 的 forward 里某个中间变量为 None权重则可能是旧版 timm 保存的新版把qkv拆成q、k、v三个独立参数。解决先创建干净环境Python 用 3.8 或 3.9torch 用 2.0 左右然后按压缩包内 requirements.txt 安装不要自己挑版本。如果源码没有 requirements我一般用pip install timm0.6.13 transformers4.30.2作为起步版本。安装后立刻跑一个最小样例随机噪声过一遍模型能输出就说明环境通过。5.2 模型输出全黑或全白训练和验证都无效现象loss 在下降但可视化输出图像要么整体接近 0全黑要么整体接近 255全白。原因数据集加载时没统一归一化。PIL 读入是 0-255 的 uint8而合成雾图脚本生成的是 0-1 浮点。如果训练输入归一化到 [0,1]标签却保持 0-255模型会一直试图“放大”输出最终饱和。解决统一在 Dataset 的__getitem__里用ToTensor()或手动x / 255.0确保输入和标签都在 [0,1]。推荐用 2.3 节的 Dataset 实现ToTensor()会自动处理。如果你顺手加了自己的预处理流水线记得在模型推理后把输出 clamp 到 [0,1] 再保存。5.3 CUDA out of memory小图也跑不动现象训练到第二个 step 报torch.cuda.OutOfMemoryError即使 batch size 已经调到 4。原因patch_size 太小或嵌入维度太高导致 token 序列过长。patch_size8 时224×224 的图会产生 28×28784 个 token每个 token 维度 192自注意力的复杂度是 O(N²)这里内存占用呈平方增长。解决第一步把 batch size 降到 1 验证模型能跑第二步用梯度累积模拟小 batch第三步如果还爆把 patch_size 改成 16 或把 depth 从 6 减到 4。对于毕设来说训练分辨率 160×160 通常也够用能显著降低显存压力。5.4 损失震荡或中途变 NaN现象前几个 epoch loss 正常下降到某一步突然变成nan或 loss 曲线像锯齿一样上下跳动。原因学习率没有预热ViT 在第一个 batch 上带着随机初始化权重梯度过大直接冲爆另外 SSIM 计算对输入范围很敏感如果输入有值超出 [0,1]内部会出现 log 负值产生 NaN。解决把优化器从普通 Adam 换成 4.2 节里的 OneCycleLR它会自动做前 10%-20% step 的线性预热训练循环里加clip_grad_norm_(model.parameters(), max_norm1.0)并在损失函数入口处对输入做pred.clamp(0, 1)。这是我在多个去雾项目里验证过的最稳妥组合。5.5 CSV 文件读取报错或列名对不上现象pandas 读 CSV 时ParserError或者pivot时提示KeyError: dim1。原因这些 CSV 有的不是标准逗号分隔可能是空格或 Tab有的读取时表头被当成数据导致列名成了第一行内容。解决读取前先强行指定分隔符和查看原始内容head -5 cifar100_vit_ti_losslandscape.csvdf pd.read_csv(cifar100_vit_ti_losslandscape.csv, sepNone, enginepython) print(df.columns)sepNone 会让 pandas 自动推断分隔符enginepython 能在处理不规则分隔时更宽容。确认列名后再执行 pivot 画图。如果列名是0,1,2这种纯数字说明文件没有表头用names[dim1, dim2, loss]手动指定即可。6. 验证与进阶拿真实雾图跑一次推理再用 PSNR 和损失景观图收尾6.1 推理脚本从本地权重到输出图训练完成后验证环节不能只看训练集上的 loss。拿一张真实雾天照片跑推理模型的表现是否自然比任何指标都直观。推理代码比训练简单得多import torch from PIL import Image import torchvision.transforms as T model.eval() img Image.open(real_haze.jpg).convert(RGB) x T.ToTensor()(img).unsqueeze(0).cuda() with torch.no_grad(): out model(x).clamp(0, 1).squeeze(0).permute(1, 2, 0).cpu().numpy() out (out * 255).astype(uint8) Image.fromarray(out).save(real_haze_removed.jpg)关键步骤就两个模型切到 eval 模式否则 BN 和 Dropout 行为不一致输出 clamp 到合法像素范围。跑通推理后如果有 paired 的测试集用 PSNR 和 SSIM 算数值指标python -m pytorch_msssim --help # 确认环境里有库算 PSNR 用skimage.metrics.peak_signal_noise_ratioSSIM 用pytorch_msssim.ssim分别计算 20 张图取平均写进文档的表格里。我自己会在答辩前一晚强制走一遍“干净环境重装 → 数据路径核对 → 30 epoch 复训 → 单张推理 → 指标计算”五步流程确保代码在导师的电脑上也能一键跑起来。6.2 两个加分项损失景观图与消融实验如果你时间充裕把第 4.3 节画出的vit_loss_landscape.png和resnet_loss_landscape.png并排放进论文“实验分析”章节说明 ViT 的收敛盆地更平坦、对权重扰动更鲁棒这比口述“Transformer 效果好”更有说服力。另一个加分项是消融实验关闭跳跃连接跑一次对比记录去掉 skip 后 PSNR 掉了多少。这份源码包的价值在于它把能出图的实验材料都备齐了剩下的就是把它们组织成一篇完整的技术报告。从那以后我每次复现类似项目都会先确认“能不能出图、能不能出表、能不能出指标”三件事再决定值不值得深挖希望帮到你。本文还有配套的精品资源点击获取