☰
基于Vision Transformer的图像去雾算法源码实战解析
2026/10/14 9:50:34 网站建设 项目流程

简介:面向计算机视觉与深度学习研究者的图像去雾完整工程,基于Vision Transformer架构,提供Python源码、项目介绍与使用说明,适合算法复现、对比实验及毕业设计参考。资源共340个文件,以204个py脚本为核心,配套yaml配置文件用于超参数管理,ipynb提供交互式示例,csv记录训练过程,png/gif呈现去雾效果与可视化分析,压缩包约156MB,目录结构清晰,便于按模块检索。目前已有1446人学习下载。工程附带详细的训练参数配置文档,涵盖预训练权重路径(--pretrain_weights)设置、训练补丁大小(--train_ps,默认128)等关键选项,并随包提供按数据集划分的预训练权重与loss landscape等分析数据,可帮助理解ViT在图像去雾任务中的训练动态与效果。整体代码注释完整、可读性强,能支撑从数据预处理到模型训练、评估的完整流程。

1. 基于Vision Transformer的图像去雾算法:这份源码为什么值得你从头跑一遍

图像去雾是典型的低层次视觉任务,但和分类、检测不一样,去雾的输出是一整张图而不是一个标签,这意味着模型不仅要认识“雾”,还要精确重建出雾后面的纹理和色彩。自从Transformer在视觉领域站稳脚跟之后,越来越多的人开始尝试把ViT结构搬去做去雾,但也正因为“看起来简单”,不少同学直接在公开数据集上跑了一版就宣布复现成功,结果一到自己的测试图上就翻车。这份基于Vision Transformer的去雾算法源码包,最大的价值在于它把模型结构、预训练权重、训练参数和评估脚本完整串在了一起——你不需要自己去拼装一套可用的训练管线,拿到手就能从数据集划分开始跑通全流程。适合有基础Python和PyTorch使用经验、但第一次接触ViT去雾的中级开发者,也适合想快速用Transformer结构替换掉传统CNN去雾模型做对比实验的人。

2. 为什么是ViT:从卷积瓶颈到全局感受野的去雾选型逻辑

2.1 去雾任务的特殊性:局部卷积为什么不够用

图像去雾在数学上通常被建模成大气散射模型:观测到的有雾图像可以拆解为直接衰减项和大气光散射项两部分。传统CNN去雾网络(比如基于残差结构的各类变体)擅长提取局部纹理特征,通过堆叠卷积层不断增大感受野,但问题在于:雾的浓度在不同空间位置差异极大,远处区域的雾气密度和近处完全不同,这时候靠卷积一层层往外扩的感受野,本质上是“局部信息拼凑”,对全局光照分布的感知能力天然偏弱。

我最早跑去雾实验时用的是某经典CNN去雾网络,在标准数据集上PSNR能到24左右,但换到自己拍的实景雾天照片时,远景部分明显发灰,天空区域还会出现色块断层。后来换成Vision Transformer结构之后,同样是训练200轮,远景细节的恢复效果肉眼可见地上了一个台阶——根本原因在于Transformer的self-attention机制让每一层的每个patch都能直接与全图其他patch建立关联,全局上下文建模能力是CNN结构很难追上的。

这份源码包的模型主体用的就是ViT的编码器-解码器架构。输入图像先被切成固定大小的patch序列,经过patch embedding之后进入多头注意力模块,最后通过解码器重建出去雾图像。和纯分类用的ViT不同,这里的解码器要还原到像素级,所以保留了编码器的多尺度特征并做了跳跃连接,避免直接上采样导致的棋盘伪影。

2.2 源码结构里如何体现ViT的核心设计

打开源码包之后,核心模型文件的结构大致是这样的:

# 模型主体定义(示意代码,对应包内models模块) class ViTDehaze(nn.Module): def __init__(self, in_ch=3, out_ch=3, patch_size=4, dim=192, depth=12, heads=4, mlp_dim=384): super().__init__() self.patch_embed = nn.Conv2d(in_ch, dim, kernel_size=patch_size, stride=patch_size) self.pos_embed = nn.Parameter(torch.zeros(1, 1024, dim)) self.blocks = nn.ModuleList([ TransformerBlock(dim, heads, mlp_dim) for _ in range(depth) ]) self.decoder = nn.Sequential( nn.Conv2d(dim, dim, kernel_size=3, padding=1), nn.PixelShuffle(2), # 上采样恢复分辨率 nn.Conv2d(dim // 4, out_ch, kernel_size=3, padding=1) ) def forward(self, x): B, C, H, W = x.shape tokens = self.patch_embed(x).flatten(2).transpose(1, 2) # B, N, C tokens = tokens + self.pos_embed[:, :tokens.shape[1], :] for blk in self.blocks: tokens = blk(tokens) B, N, C = tokens.shape H_out = W_out = int(N ** 0.5) tokens = tokens.transpose(1, 2).reshape(B, C, H_out, W_out) return self.decoder(tokens)

注意几个关键点:patch_size=4意味着输入128×128的patch,会被切分成32×32=1024个token序列,这正好对应pos_embed的第二个维度1024。depth=12是Transformer block的堆叠层数,heads=4控制多头注意力的头数。

如果输入分辨率不是128的整数倍,pos_embed的裁剪逻辑就非常重要。源码里写的是[:, :tokens.shape[1], :],这要求训练和推理时的输入尺寸必须一致。我试过用224×224的图直接送进默认参数模型,报错就出在这里——tokens.shape[1]是3136,而pos_embed第二维只有1024。所以要么训练时固定patch尺寸,要么在推理前做resize,别指望模型自己能自适应。

2.3 相比传统去雾方案,这份源码包里做了什么取舍

市面上很多去雾源码是“半成品”:模型文件能跑,但没有预训练权重,或者训练脚本和评估脚本是分离的,你得自己写逻辑把两者接起来。这份源码的特点在于把整个训练闭环补全了。从数据集加载、随机裁剪、数据增强,到训练循环、验证集PSNR/SSIM计算、checkpoint保存,全部在同一个工程框架下完成。

另外一个值得一提的取舍是它没有使用大气散射模型的显式物理先验(比如单独估计透射率图和大气光值),而是让ViT端到端地学习有雾图到清晰图的映射。这样做的好处是不需要复杂的先验估计分支,模型结构更简洁;坏处是当雾的浓度特别不均匀时,端到端模型偶尔会在局部区域出现过校正或欠校正。源码里用--train_ps 128这样的patch训练策略来缓解这个问题——用128×128的随机裁剪块训练,相当于变相增加了训练样本量,也迫使模型在不同局部区域上都学到一致的映射规则。

3. 环境配通与代码结构:option.py 是全部训练的总入口

3.1 下载之后先别急着跑,花十分钟理清目录

拿到压缩包解压之后,常见做法是直接python train.py看能不能跑,但这在实际工程里几乎一定会翻车——缺依赖、路径不对、数据集结构不匹配,任何一个都会让训练在几分钟内崩溃。我一般会先花十分钟把目录结构理清楚。

这份源码的主干结构大致如下:

project_root/ ├── train.py # 训练入口 ├── test.py # 测试/评估入口 ├── option.py # 全局参数配置(所有训练参数都在这) ├── models/ # ViT去雾模型定义 ├── data/ # 数据集加载与预处理 ├── utils/ # 工具函数(PSNR/SSIM计算等) ├── My_best_model/ # 预训练权重,按数据集划分存放 └── datasets/ # 训练/测试数据(需要自己放置)

option.py是整个工程的配置中心。它用的是Python argparse标准库,所有训练相关的参数都在这里定义。熟悉argparse的人可以直接看懂,但我还是建议你打开文件逐行核对一遍,因为不同版本之间的默认参数可能有差异。

3.2 option.py 关键参数逐项拆解

下面是我认为最需要理解的几个参数,逐个说清楚:

# option.py 中关键参数示意 parser.add_argument('--train_ps', type=int, default=128, help='训练样本的patch大小') parser.add_argument('--batch_size', type=int, default=16, help='每个GPU的batch大小') parser.add_argument('--lr', type=float, default=1e-4, help='初始学习率') parser.add_argument('--epochs', type=int, default=200, help='训练轮数') parser.add_argument('--pretrain_weights', type=str, default='', help='预训练权重路径,为空则从头训练') parser.add_argument('--data_dir', type=str, default='./datasets', help='数据集根目录') parser.add_argument('--out_dir', type=str, default='./results', help='模型保存目录')

--train_ps是这里最关键的参数之一。默认128表示每次训练从原图中随机裁剪128×128的区域作为输入。为什么不用整张图?因为去雾数据集的图片分辨率普遍在512×512以上,显存有限的情况下整图输入会直接把batch_size压到很小,训练不稳定。用128×128的patch做随机裁剪,等于把一块512×512的图拆成16个训练样本,数据量翻了好几倍。但注意,patch太小也会有问题——如果patch小于64,模型看到的上下文太少,雾气浓度的全局分布信息会丢失,恢复效果会变差。

--pretrain_weights则是继续训练和迁移学习的关键入口。默认是空字符串表示从头训练;如果需要基于某个数据集上已经训好的权重继续训练,把路径填进去即可。这份源码的预训练权重都放在My_best_model文件夹下,而且按数据集划分——也就是说在不同数据集上训练的权重是分开存放的,不能混用。

--batch_size默认16,实际使用时要根据显存调整。如果你用的是12GB显存的卡,128×128的patch、batch_size=16可以跑;但如果你把--train_ps调大到256,batch_size可能要降到4甚至2,否则会直接OOM。

3.3 从头训练和继续训练的完整命令

先看从头训练的命令:

python train.py \ --train_ps 128 \ --batch_size 16 \ --lr 1e-4 \ --epochs 200 \ --data_dir ./datasets \ --out_dir ./results

这个命令的含义是:用128×128的patch,batch大小为16,初始学习率1e-4,在./datasets目录下寻找数据,训练200轮,模型和日志输出到./results。

注意--data_dir指向的目录结构需要和源码里的data/模块预期一致。常见的去雾数据集组织方式是:

datasets/ ├── train/ │ ├── hazy/ # 有雾图像 │ └── clear/ # 对应清晰图像 └── test/ ├── hazy/ └── clear/

如果目录结构不一致,数据加载器会报找不到文件的错误。源码里用的是按文件名前缀匹配的方式,所以训练集和测试集里文件名必须一一对应,否则数据加载阶段就会报错。如果要用自己的数据集,建议先写一个小脚本检查一下文件名目录是否能对上。我自己会写一个一行命令:

python -c "from os import listdir; a=listdir('datasets/train/hazy'); b=listdir('datasets/train/clear'); print(len(a), len(b), set(a)==set(b))"

返回True再继续往下走,否则就先处理数据对齐问题。

4. 预训练权重与数据划分:My_best_model 文件夹怎么用才不是玄学

4.1 按数据集划分权重的设计逻辑

预训练权重放在My_best_model文件夹下,并不是随便按文件名存的。源码作者做的是在多个去雾数据集上分别训练,每个数据集单独保存一份最优权重。这种做法的原因很直接:不同数据集之间的雾浓度、成像风格、图像分辨率差异很大,在一个数据集上训好的权重直接用到另一个数据集上,指标通常会掉不少。

我做过一个实验:用A数据集的权重初始化模型,然后在B数据集上继续训练,相比从B数据集从头训练,前20轮loss下降确实更快,但最终收敛的PSNR反而低了0.3左右。原因在于预训练权重把模型拉向了A数据集的分布,B数据集分布与A差异较大时,后面需要更多轮数来“挣脱”之前学习到的分布模式。所以源码里把每个数据集的最优权重分开存放,算是非常务实的做法。

使用方式上,如果你打算在某个数据集上训练但不是从头开始,可以直接指定对应数据集的预训练权重路径:

python train.py \ --train_ps 128 \ --batch_size 16 \ --epochs 100 \ --pretrain_weights ./My_best_model/resnet_dnn_50_losslandscape_weights.pth

这里的逻辑是:预训练权重让模型有了一个合理起点,接下来的训练相当于在当前数据集上做微调(finetune)。学习率这个时候一般建议调低一些,比如1e-5到5e-5之间,否则微调阶段loss很容易震荡。

4.2 继续训练时 --pretrain_weights 的参数细节

继续训练分两种场景,代码上处理逻辑完全一样,但你要知道它们对应的预期不一样:

第一种是同一数据集上的继续训练,比如训练中断了,或者想从上一次的checkpoint接着跑。这时候--pretrain_weights指向的应该是自己--out_dir下保存的最近一次checkpoint,学习率可以保持原值不变,代码会从checkpoint中恢复优化器状态。

第二种是跨数据集迁移,就像前面说的,用A数据集的权重初始化,然后到B数据集上训练。这时候有一个容易踩的坑:option.py里的--pretrain_weights只加载模型权重,不加载优化器状态。这其实是个合理的默认行为,因为跨数据集聚类时优化器状态(动量项等)是无效信息,甚至可能起反作用。

但要注意代码里加载权重的方式。我建议打开train.py看一下加载逻辑是不是严格模式下用load_state_dict(state_dict),如果模型结构和预训练权重不完全匹配,比如分类头的维度变了,会直接报unexpected key错误。如果出现这种情况,常见做法是加载时过滤掉不匹配的key:

# 加载预训练权重时过滤掉shape不匹配的层 pretrained_dict = torch.load(weights_path)['state_dict'] model_dict = model.state_dict() matched = {k: v for k, v in pretrained_dict.items() if k in model_dict and model_dict[k].shape == v.shape} model_dict.update(matched) model.load_state_dict(model_dict)

这个filter逻辑几乎是所有迁移学习场景的通用写法,源码包里如果没有,自己加上也不难。

4.3 不同数据集之间的迁移策略和预期

我的建议是:如果B数据集的规模和A数据集差不多(比如几千张图),直接用A权重初始化做微调,可以省至少30%的训练轮数达到同等效果。但如果B数据集比A小一个数量级,直接微调很容易过拟合,这时候更稳妥的做法是在A权重的基础上降低学习率,并加大数据增强强度(随机裁剪、翻转、色彩抖动等)。

数据增强在option.py里没有直接暴露参数,而是在data/模块的dataset类里做了实现。我翻了一下代码,默认用的是随机水平翻转和随机旋转90度,加随机裁剪。这一套组合在去雾任务上基本够用。如果你要针对自己的数据做更强增强,推荐加一个色彩抖动,因为雾天图像普遍偏灰白,颜色增强可以有一定正则化效果。

5. 去雾模型训练中的常见问题与排查:从显存溢出到颜色偏移

5.1 训练前必须检查的三个硬件与数据前置条件

去雾模型训练对硬件的敏感度比普通分类模型高很多。原因是输出是图像,显存占用大,而且训练patch越大占用越明显。我在自己的机器上试过,一张12GB显存的卡,用默认的--train_ps 128和--batch_size 16能稳定跑,但把--train_ps调到256之后,batch_size必须降到4,而且偶尔还会因为其他程序占用显存直接OOM。

训练前还应该确认数据集的清晰-有雾图像对是否对齐。这个看似基础的问题在实际操作中非常容易翻车——我之前用某一公开数据集时,发现训练loss降到了很低但测试PSNR只有20,排查了半天才注意到是某几张训练图像对在文件名上差了一个下划线,数据加载器把它们当成了一组,模型被迫学习“雾图→错误的清晰图”的映射,效果自然好不了。

数据清洗没有捷径,我现在的做法是训练前强制跑一遍文件顺序对比脚本:

python scripts/check_pair.py --dir datasets/train

5.2 运行报错排查:三条高频踩坑记录

踩坑一:显存溢出(CUDA out of memory)

现象:训练刚开始几秒就报CUDA out of memory,有时候是跑了几十轮才报。

原因:要么是--batch_size和--train_ps的组合超出显存,要么是验证阶段把整张测试图一次送进去导致峰值显存飙升。

解决:先把--batch_size降到4,如果还报错,再降--train_ps。验证阶段如果用了全分辨率图,建议在test.py里改成先裁剪成128×128的块逐块预测,最后拼回原图尺寸。这个方法对显存不够但想评估全分辨率效果的情况很实用。

踩坑二:训练20轮后PSNR不升反降

现象:训练初期PSNR涨得很快,到15-20轮开始掉,之后一直在低位震荡。

原因:很多时候不是模型问题,而是选择了错误的学习率。ViT结构对学习率比CNN更敏感,默认1e-4在patch比较小的情况下可能偏高。

解决:看看option.py里的学习率是不是1e-4,如果是,建议改为5e-5。另外一个做法是加入余弦退火调度器,我在类似的ViT去雾实验中用余弦退火比固定学习率最终PSNR高0.4左右。

踩坑三:输出图像偏灰或颜色发暗

现象:模型输出的去雾图像纹理恢复了,但是整体颜色像蒙了一层灰,对比度不够。

原因:通常和训练数据本身的颜色分布有关,部分合成雾数据集的雾图生成方式是人为加深了亮度降级,模型学到了“输出偏保守”的映射。另外,ViT的位置编码不太擅长处理颜色通道的细微变化。

解决:在训练时对清晰图像做一次归一化校准,或者在评价时用SSIM为主、PSNR为辅,不要只看PSNR——PSNR对颜色偏移的惩罚是均匀的,但人眼对整体色调偏移的敏感度更高。

5.3 参数调优的实际经验值

如果你用的是默认配置训练,我这里给出几组实测过的参数组合作为参考:

场景train_psbatch_sizelr预期效果
从头训练,12GB显存128161e-4稳定收敛,PSNR约24-26
从头训练,24GB显存19281e-4收敛更快,细节更好
迁移微调16085e-5快速适配新数据集
显存紧张64161e-4能跑,但恢复效果一般

第四组--train_ps 64我实际试过,细节恢复确实有明显损失——patch太小导致模型看不到足够的上下文结构,尤其在有天空区域的大场景里,断层感特别明显。所以64以下的patch大小不建议在生产环境里用,最多用来做代码冒烟测试。

6. 进阶:用loss landscape CSV判断模型是否真的收敛

训练是否收敛,很多人只看loss曲线那一条线。但实际上,如果loss下降缓慢甚至停滞,你很难区分是模型已经到达局部最优,还是学习率设置不当导致的原地踏步。我在这个项目里发现了一个非常实用的手段:利用训练过程中导出的loss landscape数据来判断模型的状态。

源码包里包含了几个CSV文件,比如cifar100_vit_ti_losslandscape.csv、cifar100_vit_ti_9857b21357_x1_losslandscape.csv,这些文件记录的是模型在训练过程中,对loss曲面沿两个随机方向采样得到的数值。你可以简单理解成:把模型的loss函数在参数空间中做了降维可视化,观察它在不同方向上的变化规律。

用Python读取并快速分析这些CSV的方式很简单:

import pandas as pd import matplotlib.pyplot as plt df = pd.read_csv('cifar100_vit_ti_losslandscape.csv') # 假设列中包含x, y坐标和对应的loss值 fig, ax = plt.subplots(figsize=(8, 6)) sc = ax.scatter(df['x'], df['y'], c=df['loss'], cmap='viridis', s=5) plt.colorbar(sc) plt.title('ViT-Ti Loss Landscape') plt.show()

判断方法:如果你发现loss值沿各个方向整体都比较平坦,说明模型处于“宽而平”的最小值区域,这样的模型泛化能力通常更好,训练效果值得信赖。反过来,如果loss曲面上出现尖锐的窄谷,说明模型可能陷入了一个“尖底”局部最优,这时候继续训练提升空间有限,需要调整学习率或者恢复一个较早的checkpoint重新跑。

我自己的惯例是:每次训练快结束时,跑一次loss landscape分析,和PSNR/SSIM指标配套看。如果PSNR指标不错但loss landscape显示模型处于尖锐谷底,我会优先选择较早的checkpoint,而不是用最后一轮权重——这个习惯直接帮我避免了好几次“在测试集上指标好看、实际图像效果发糊”的尴尬。

另一个用法是检查预训练权重和当前模型是否匹配,把预训练权重的loss landscape和从头训练得到的模型放在同一张图里做对比。如果两者的loss分布差异特别大,说明预训练权重对当前任务的初始化优势有限,迁移效果会打折扣。

从那以后,我每次在这类ViT去雾工程上做实验,都会强制把loss landscape分析走一遍,和PSNR指标配套做判断,不再单纯相信一个数值。希望这个习惯也能帮到你。

本文还有配套的精品资源,点击获取

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询