简介图像去雾是计算机视觉中底层图像复原的关键任务旨在从受大气散射影响的退化图像中恢复清晰场景。其核心原理在于建模雾的全局物理退化过程传统方法依赖物理先验而深度学习方法通过数据驱动学习雾到清晰图像的映射。Vision Transformer凭借其自注意力机制能有效建模图像全局依赖关系为处理雾的全局分布特性提供了新的技术路径在图像复原领域展现出重要价值。本文聚焦于将ViT架构应用于真实场景图像去雾通过构建编码器-解码器结构、设计多损失函数并利用NH-HAZE、NTIRE2019等多个真实雾霾数据集进行联合训练与优化详细阐述了模型构建、训练策略及调优实践为深度学习在底层视觉任务中的应用提供了完整范例。1. 项目概述当ViT遇上图像去雾最近在整理过往的项目资料翻到了一个挺有意思的“老”项目——基于Vision TransformerViT架构的图像去雾算法研究。说它“老”是因为ViT在视觉领域的应用热潮已经持续了好几年但说它“新”是因为将纯Transformer架构系统性地应用于真实雾霾图像去雾并覆盖NH-HAZE、NTIRE2019、I-HAZE、O-HAZE这四种主流真实数据集进行训练和测试在当时乃至现在都是一个兼具挑战性和实用性的课题。这个项目打包文件.zip里就包含了从模型构建、数据集处理到训练测试的全套代码与实验记录。图像去雾简单说就是让雾天拍的照片变清晰。这听起来像是手机里的一个滤镜功能但其背后的技术深度远超想象。传统的去雾方法大多基于大气散射物理模型进行先验估计效果在复杂多变的真实场景中常常不尽如人意。深度学习的介入尤其是卷积神经网络CNN让数据驱动的去雾成为了主流。然而CNN在捕捉长距离依赖关系上存在天然局限而雾的分布恰恰可能是全局性的、非局部的。这正是Vision Transformer的用武之地。ViT通过自注意力机制能够建模图像所有区块patch之间的全局关系理论上更适合理解整幅图像的雾浓度分布与结构信息。这个项目的核心就是探索这种“理论上”的优势能否在NH-HAZE近红外-可见光配对、NTIRE2019高分辨率真实雾图、I-HAZE室内雾和O-HAZE室外雾这四种特性各异的真实雾霾数据集上转化为实际更优的去雾效果。它不是一个简单的模型套用而是一次针对特定任务的深度架构适配与优化实践。无论你是刚入门深度学习想找一个有深度的实战项目还是已经在计算机视觉领域工作、希望深入了解Transformer在底层视觉任务中的应用这个项目都能提供从理论到代码的完整视角。2. 核心思路与方案选型为什么是ViT在动手写代码之前我们需要想清楚一个根本问题在图像去雾这个任务上为什么选择Vision Transformer而不是继续优化更成熟的CNN架构比如各种U-Net变体这决定了我们整个项目的技术路线。2.1 从CNN的局限到ViT的潜力卷积神经网络的成功建立在两个核心假设上局部连通性和平移不变性。这对于提取边缘、纹理等局部特征非常有效。在去雾任务中CNN可以很好地恢复局部细节。但是雾的物理形成过程是全局性的。大气中的悬浮颗粒对光线的散射作用导致图像对比度下降、颜色偏移这种退化效应并非只影响图像的某个小区域而是与整个场景的深度、光照和雾浓度全局相关。一个区域的最佳去雾强度可能需要参考图像另一端区域的信息来判断。CNN通过堆叠卷积层来扩大感受野以间接获取全局信息但这是一种效率不高的方式且深层网络容易带来梯度消失或爆炸等问题。而Transformer的自注意力机制允许模型在第一步就直接计算图像中任意两个区块之间的关系权重。这意味着模型可以“一眼”就看到天空区域的浓厚雾气与地面景物微弱细节之间的关联从而做出更协调、全局一致的去雾决策。这是选择ViT架构最根本的动机。2.2 针对去雾任务的ViT架构适配标准的ViT是为图像分类设计的将图像分割成固定大小的区块线性投影为序列加上位置编码然后送入Transformer编码器最后用一个分类头输出类别。显然去雾是一个像素级的回归任务输出清晰图像我们需要对标准ViT进行关键改造编码器-解码器结构我们采用类似U-Net的对称结构但主干是Transformer。编码器部分由多个Transformer块组成用于提取多尺度的全局上下文特征。解码器部分则需要将序列特征重新映射回二维图像空间。这里通常采用“上采样Transformer块”或结合轻量级卷积进行特征重组和上采样。局部-全局特征融合纯自注意力机制对局部细节的捕捉有时不如卷积直接。为了兼顾两者一种常见的策略是采用混合架构Hybrid Architecture即在输入阶段先使用一个小型CNN如几个卷积层提取初步的局部特征图再将特征图分割成区块送入Transformer。或者在Transformer块中引入卷积前馈网络或使用窗口注意力如Swin Transformer的思想来增强局部建模能力。跳跃连接为了缓解训练难度并保留更多底层细节这对恢复清晰纹理至关重要必须在编码器和解码器的对应层级之间添加跳跃连接。这与U-Net的思想一脉相承确保解码器在上采样过程中能融合编码器捕获的高分辨率细节。2.3 多数据集联合训练的考量项目包含了四个真实雾霾数据集NH-HAZE、NTIRE2019、I-HAZE、O-HAZE。它们各有特点NH-HAZE包含近红外和可见光配对图像提供了多模态信息挑战在于如何有效融合近红外通道受雾影响较小来辅助可见光去雾。NTIRE2019分辨率高雾浓度多样场景复杂是检验模型泛化能力的试金石。I-HAZE室内场景雾源可控但受室内灯光和复杂反射影响颜色保真度挑战大。O-HAZE室外自然场景雾浓度不均匀包含丰富的自然景物细节。将它们联合训练而非单独训练四个模型有几个好处第一大幅增加训练数据的多样性和数量有助于模型学习更通用的去雾先验避免过拟合到某种特定雾类型。第二模型被迫适应不同成像条件室内/室外、不同传感器其鲁棒性会更强。当然挑战在于数据分布差异可能带来优化冲突需要精心设计数据加载策略和损失函数。注意联合训练时最简单的策略是混合所有数据。但更精细的做法可以是“课程学习”或“批次内平衡”例如确保每个训练批次都包含来自不同数据集的样本防止模型被数据量最大的那个数据集主导。3. 项目实战从数据准备到模型训练理论说得再多不如一行代码。接下来我们深入这个项目的实战环节看看如何一步步搭建并训练这个ViT去雾模型。3.1 真实雾霾数据集处理全攻略拿到NH-HAZE、NTIRE2019等数据集后第一步不是直接开训而是系统的数据预处理与增强。这一步的质量直接决定了模型性能的上限。数据解压与结构整理通常每个数据集都有其官方结构。我们需要编写统一的脚本将它们整理成一种固定的格式。我推荐的目录结构如下dataset_root/ ├── train/ │ ├── hazy/ # 存放所有训练集的雾图 │ │ ├── NH_HAZE_1_hazy.png │ │ ├── NTIRE_2019_1_hazy.png │ │ └── ... │ └── gt/ # 存放所有训练集的清晰图ground truth │ ├── NH_HAZE_1_gt.png │ ├── NTIRE_2019_1_gt.png │ └── ... └── val/ (或 test/) ├── hazy/ └── gt/这样做的好处是数据加载器可以无视图像来源统一读取。你需要根据每个数据集的说明文件将雾图hazy和清晰图gt正确配对并移动到对应目录。特别注意文件名的一致性以便于配对。关键预处理步骤图像对齐检查对于真实数据集尤其是室外拍摄的尽管声称是配对的但极轻微的抖动或透视变化都可能存在。务必进行抽查可以通过计算互信息或简单的像素差来肉眼检查。如果发现严重不对齐该样本可能需要剔除或使用图像配准算法进行校正否则模型永远学不到正确的映射。分辨率统一与区块化ViT要求输入固定尺寸。我们需要将图像裁剪或缩放到统一大小如256x256或512x512。对于训练更常用的做法是随机裁剪Random Crop这本身就是一种数据增强。对于高分辨率的NTIRE2019图像随机裁剪能产生大量训练样本。同时要确保裁剪对雾图和清晰图在同一位置进行。数据增强策略除了随机裁剪针对去雾任务我强烈推荐使用随机水平/垂直翻转简单有效。颜色抖动轻微调整亮度、对比度和饱和度。因为雾会影响颜色适度的颜色扰动可以让模型对颜色变化更鲁棒但强度不宜过大以免破坏颜色恒常性约束。归一化将像素值从[0, 255]缩放到[0, 1]或[-1, 1]。同时计算整个训练集的均值和标准差用于标准化输入加速模型收敛。针对NH-HAZE的特殊处理NH-HAZE提供了近红外NIR通道。一个有效的做法是将三通道的RGB雾图与单通道的NIR图像拼接成一个4通道的输入。在模型的第一层可以使用一个4通道输入的卷积层来融合这些信息。清晰图GT仍然是3通道RGB。这要求我们对数据加载器进行修改使其能同时加载RGB和NIR图像。3.2 Vision Transformer去雾模型构建详解这里我们设计一个相对简洁但有效的ViT去雾网络。我们称之为“DehazeViT”。以下是用PyTorch框架搭建的核心代码解析import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange, repeat from einops.layers.torch import Rearrange # 核心组件Transformer块 class TransformerBlock(nn.Module): def __init__(self, dim, heads, dim_head, mlp_dim, dropout0.): super().__init__() self.norm1 nn.LayerNorm(dim) self.attn nn.MultiheadAttention(dim, heads, dropoutdropout, batch_firstTrue) # PyTorch原生Attention self.norm2 nn.LayerNorm(dim) self.mlp nn.Sequential( nn.Linear(dim, mlp_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(mlp_dim, dim), nn.Dropout(dropout) ) def forward(self, x): # x: [batch_size, num_patches, dim] x_norm self.norm1(x) attn_output, _ self.attn(x_norm, x_norm, x_norm) x x attn_output x_norm self.norm2(x) mlp_output self.mlp(x_norm) x x mlp_output return x # 编码器部分由多个Transformer块堆叠而成 class TransformerEncoder(nn.Module): def __init__(self, depth, dim, heads, dim_head, mlp_dim, dropout0.): super().__init__() self.layers nn.ModuleList([]) for _ in range(depth): self.layers.append(TransformerBlock(dim, heads, dim_head, mlp_dim, dropout)) def forward(self, x): for layer in self.layers: x layer(x) return x # 主模型DehazeViT class DehazeViT(nn.Module): def __init__(self, image_size256, patch_size16, in_channels3, out_channels3, dim768, depth12, heads12, mlp_dim3072, dropout0.1): super().__init__() assert image_size % patch_size 0, 图像尺寸必须能被区块大小整除 num_patches (image_size // patch_size) ** 2 patch_dim in_channels * patch_size * patch_size # 1. 区块嵌入与位置编码 self.to_patch_embedding nn.Sequential( Rearrange(b c (h p1) (w p2) - b (h w) (p1 p2 c), p1patch_size, p2patch_size), nn.Linear(patch_dim, dim), ) self.pos_embedding nn.Parameter(torch.randn(1, num_patches 1, dim)) # 1 for cls token self.cls_token nn.Parameter(torch.randn(1, 1, dim)) # 2. Transformer编码器 self.transformer TransformerEncoder(depth, dim, heads, dim//heads, mlp_dim, dropout) # 3. 解码器这里简化为一个线性层映射回像素空间实际可替换为更复杂的上采样网络 self.decoder nn.Sequential( nn.LayerNorm(dim), nn.Linear(dim, patch_dim), Rearrange(b (h w) (p1 p2 c) - b c (h p1) (w p2), himage_size//patch_size, wimage_size//patch_size, p1patch_size, p2patch_size, cout_channels), ) def forward(self, img): # img: [B, C, H, W] x self.to_patch_embedding(img) # [B, num_patches, dim] b, n, _ x.shape cls_tokens repeat(self.cls_token, 1 1 d - b 1 d, bb) x torch.cat((cls_tokens, x), dim1) # 拼接cls token x self.pos_embedding[:, :(n 1)] x self.transformer(x) # 丢弃cls token只处理图像区块特征 x x[:, 1:, :] out self.decoder(x) return out关键点解析区块嵌入使用einops库的Rearrange操作将图像[B, C, H, W]重排为[B, (H*W)/(P*P), P*P*C]再通过线性层投影到固定维度dim。这是将2D图像转换为1D序列的标准操作。位置编码由于Transformer本身没有位置概念必须添加可学习的位置编码pos_embedding让模型知道每个区块在原始图像中的空间位置。这是至关重要的。分类令牌cls_token是ViT用于分类的在去雾任务中我们可以选择保留它作为全局特征的聚合也可以在解码前丢弃。上面的代码在解码时丢弃了它专注于图像区块特征。解码器简化上述代码的解码器过于简单一个线性层。在实际项目中你需要一个更强的解码器。一个有效的方案是使用多个“上采样Transformer块”或一个轻量的CNN上采样网络例如几个转置卷积或PixelShuffle层。解码器的输入是Transformer编码器输出的序列特征需要将其重塑并上采样回原始图像尺寸。3.3 损失函数设计与训练策略去雾是一个回归任务损失函数的设计直接影响复原图像的质量。单一损失函数往往不够需要组合拳。1. 像素级损失最基础的是L1损失平均绝对误差MAE。相比L2损失MSEL1损失对异常值不那么敏感产生的图像通常更清晰边缘更锐利是底层视觉任务的常用选择。loss_l1 nn.L1Loss()(pred_image, gt_image)2. 感知损失仅靠像素匹配容易导致结果过于平滑丢失纹理。感知损失通过一个预训练的网络如VGG16来比较预测图像和真实图像在特征空间的距离鼓励它们在高级语义特征上相似从而保留更自然的纹理和结构。# 假设vgg是一个预训练VGG的特定层特征提取器 loss_perceptual F.l1_loss(vgg(pred_image), vgg(gt_image))3. 对抗损失如果想追求极致的视觉真实感可以引入生成对抗网络GAN的思想。添加一个判别器网络试图区分去雾图像和真实清晰图像。生成器我们的DehazeViT则试图“欺骗”判别器。对抗损失能有效提升结果的视觉质量使图像看起来更自然但训练更不稳定。# 简化示意实际需定义判别器D loss_adv -torch.mean(D(pred_image)) # 对于生成器希望判别器给高分4. 总损失通常采用加权和。total_loss lambda_l1 * loss_l1 lambda_perceptual * loss_perceptual lambda_adv * loss_adv在项目初期建议从lambda_l11.0,lambda_perceptual0.1,lambda_adv0.0开始先稳定训练再逐步引入其他损失。训练策略与超参数选择优化器AdamW优化器是目前训练Transformer的首选它修正了Adam的权重衰减方式通常能带来更好的泛化性能。初始学习率可以设为3e-4或5e-4。学习率调度使用余弦退火调度Cosine Annealing LR Scheduler或带热重启的余弦退火让学习率从初始值平滑下降至0有助于模型收敛到更好的局部最优。批量大小在显存允许的情况下尽可能大。ViT模型对批量大小比较敏感较大的批量有助于稳定训练。如果显存不足可以使用梯度累积来模拟大批量。训练轮数由于是多数据集联合训练数据量较大通常需要训练至少100-200个epoch。密切监控验证集上的损失和指标如PSNR, SSIM当指标在连续多个epoch不再提升时可以提前停止。实操心得训练初期损失可能震荡较大。一个有用的技巧是使用梯度裁剪torch.nn.utils.clip_grad_norm_将梯度范数限制在一个阈值内如1.0这能有效防止训练初期因不稳定而崩溃。另外在引入对抗损失时可以先单独用L1损失训练几十个epoch得到一个不错的预训练模型然后再开启对抗训练这样会稳定得多。4. 多数据集训练技巧与评估指标联合训练四个数据集并非简单混合其中有许多技巧可以提升最终模型的鲁棒性和性能。4.1 批次内数据平衡与采样策略如果简单地将所有数据随机混合数据量大的数据集如NTIRE2019会主导训练过程模型可能会在它上面过拟合而在小数据集如I-HAZE上表现不佳。为了解决这个问题可以采用批次内平衡采样。具体实现是在每一个训练批次batch中确保包含来自每个数据集的样本。例如如果批次大小为16我们可以设定每个数据集贡献4张图像。这需要自定义PyTorch的Dataset和Sampler。from torch.utils.data import Dataset, DataLoader, ConcatDataset from itertools import cycle, islice class BalancedBatchSampler: def __init__(self, datasets, batch_size, samples_per_dataset): self.batch_size batch_size self.samples_per_dataset samples_per_dataset # 每个数据集每批贡献数 self.dataset_indices [] for i, dataset in enumerate(datasets): indices list(range(len(dataset))) # 打乱每个数据集自身的索引 np.random.shuffle(indices) # 为每个索引标记其所属的数据集ID self.dataset_indices.extend([(i, idx) for idx in indices]) # 打乱所有数据 np.random.shuffle(self.dataset_indices) def __iter__(self): # 这是一个简化的平衡采样逻辑实际实现可能需要更复杂的循环 batch [] count_per_ds_in_batch {i:0 for i in range(len(datasets))} for ds_id, idx in self.dataset_indices: if count_per_ds_in_batch[ds_id] self.samples_per_dataset: batch.append((ds_id, idx)) count_per_ds_in_batch[ds_id] 1 if len(batch) self.batch_size: yield batch batch [] count_per_ds_in_batch {i:0 for i in range(len(datasets))} if batch: yield batch # 使用 # train_datasets [dataset_NH, dataset_NT, dataset_I, dataset_O] # sampler BalancedBatchSampler(train_datasets, batch_size16, samples_per_dataset4) # loader DataLoader(ConcatDataset(train_datasets), batch_samplersampler)4.2 客观评估指标PSNR与SSIM模型训练好后我们需要定量评估其去雾效果。最常用的两个全参考图像质量评估指标是峰值信噪比衡量去雾图像与清晰图像之间的像素级误差值越高代表失真越小。计算简单但与人眼视觉感受不完全一致。import torch def calculate_psnr(img1, img2, max_val1.0): mse torch.mean((img1 - img2) ** 2) if mse 0: return float(inf) return 20 * torch.log10(max_val / torch.sqrt(mse))结构相似性指数从亮度、对比度、结构三个方面比较图像相似度取值范围[-1, 1]值越接近1越好更符合人眼主观感受。通常使用滑动窗口计算。# 可以使用现成的库如piq或pytorch-msssim from piq import ssim ssim_index ssim(pred_image, gt_image, data_range1.0)在测试时我们应在每个数据集独立的测试集上计算平均PSNR和SSIM。一个强大的模型应该在四个数据集上都取得相对均衡且较高的指标。如果某个数据集指标明显偏低说明模型对其泛化能力不足可能需要调整训练策略或数据增强。4.3 主观视觉评估与案例分析定量指标重要但“看起来好不好”同样关键。在项目报告中必须包含丰富的视觉对比图。将输入雾图、我们的DehazeViT输出、其他经典方法如DCP、AOD-Net、FFA-Net等的输出以及真实清晰图GT放在一起对比。分析要点去雾彻底性浓雾区域是否被有效移除远处的景物是否变得清晰颜色保真度去雾后的颜色是否自然有无明显的颜色失真或偏色例如整体发蓝或发黄这在I-HAZE室内数据集中尤其需要注意。细节恢复纹理、边缘是否清晰有无过度平滑或产生伪影如光晕、块效应整体协调性图像整体亮度、对比度是否自然有无局部过亮或过暗通过案例分析我们能更直观地发现模型的优缺点。例如模型可能在薄雾场景表现优异但在处理浓雾与天空交界处时产生光晕或者对室内人造光源下的雾处理不佳。这些观察能为后续模型改进提供明确方向。5. 常见问题、调优与避坑指南在实际操作这个项目的过程中我遇到了不少坑也总结了一些调优经验。5.1 训练不稳定与发散问题现象损失值变成NaN或者PSNR/SSIM在训练过程中骤降。排查与解决检查数据首先确认输入数据是否包含异常值如NaN或Inf。确保数据加载和预处理流程正确图像像素值在归一化后处于合理范围如[0,1]。降低学习率这是最常见的原因。尝试将学习率降低一个数量级例如从3e-4降到3e-5。启用梯度裁剪在优化器更新权重之前添加torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)。检查损失函数权重如果使用了对抗损失其权重lambda_adv可能设置过高。尝试将其设为0先只用L1损失训练稳定后再慢慢增加。简化模型如果使用的是非常深的Transformer尝试减少深度depth或注意力头数heads先从小模型开始训练成功再逐步放大。5.2 模型过拟合与泛化能力差问题现象在训练集上PSNR很高但在验证集或未见过的测试集上指标很低。排查与解决加强数据增强增加更多样化的数据增强如随机旋转小角度、添加高斯噪声、模拟不同浓度的雾通过调整大气光A和透射率t来合成雾图等。对于ViT甚至可以对区块进行随机掩码Masking作为正则化手段。使用权重衰减与Dropout确保AdamW优化器中的权重衰减参数weight_decay已设置如0.05。在Transformer块和前馈网络中使用Dropout如dropout0.1。早停持续监控验证集损失当其在连续10-20个epoch内不再下降时果断停止训练。模型集成训练多个不同初始化或不同超参数的模型对它们的预测结果进行平均可以有效提升泛化性能。5.3 显存不足与训练速度慢问题现象GPU显存溢出OOM或每个epoch耗时过长。排查与解决减小输入尺寸或批次大小这是最直接的方法。将图像裁剪尺寸从512x512降到256x256。使用混合精度训练PyTorch的torch.cuda.amp模块可以自动使用FP16进行前向和反向传播显著减少显存占用并加速训练。需注意尺度缩放Grad Scaler以避免下溢。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() with autocast(): output model(input) loss criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()梯度累积如果目标批次大小是16但显存只够放4张图可以设置梯度累积步数为4。每4个前向-反向传播步骤才真正更新一次权重optimizer.step()和optimizer.zero_grad()。检查点加载与保存只保存模型的状态字典state_dict而不是整个模型对象。定期保存检查点以便从中断处恢复。5.4 去雾结果存在特定瑕疵问题现象去雾后图像出现颜色失真、光晕、或局部区域模糊。排查与解决颜色失真检查感知损失是否权重过大或使用的VGG层太深倾向于语义特征而非颜色。可以尝试在较浅的VGG层如relu1_2,relu2_2计算感知损失。同时在损失函数中加入颜色损失例如在Lab颜色空间计算L1损失因为Lab空间更符合人眼对颜色的感知。光晕效应通常发生在景深突变或物体边缘。这可能是模型过度平滑或感受野有限导致的。可以尝试在损失函数中加入梯度损失对预测图像和真实图像的梯度图计算L1损失以强化边缘保持。def gradient_loss(pred, gt): pred_gray 0.299*pred[:,0,:,:] 0.587*pred[:,1,:,:] 0.114*pred[:,2,:,:] gt_gray 0.299*gt[:,0,:,:] 0.587*gt[:,1,:,:] 0.114*gt[:,2,:,:] pred_grad_x torch.abs(pred_gray[:, :, 1:] - pred_gray[:, :, :-1]) pred_grad_y torch.abs(pred_gray[:, 1:, :] - pred_gray[:, :-1, :]) gt_grad_x torch.abs(gt_gray[:, :, 1:] - gt_gray[:, :, :-1]) gt_grad_y torch.abs(gt_gray[:, 1:, :] - gt_gray[:, :-1, :]) loss F.l1_loss(pred_grad_x, gt_grad_x) F.l1_loss(pred_grad_y, gt_grad_y) return loss局部模糊模型可能过于依赖全局信息而忽略了局部细节。考虑在架构中引入局部注意力机制或卷积旁路。例如可以使用Swin Transformer的窗口注意力或者在编码器每个阶段后将Transformer特征与一个浅层CNN提取的细节特征进行融合。这个基于Vision Transformer的图像去雾项目就像一次深入的探险。它不仅仅是将一个热门架构应用于新任务更要求你深入理解去雾任务的物理本质与数据特性并针对性地调整模型、损失和训练策略。从数据预处理的小心对齐到模型设计中全局与局部权衡再到训练时各种“炼丹”技巧每一步都充满了挑战和收获。最终当你看到模型成功驱散浓雾还原出清晰生动的图像时那种成就感是实实在在的。希望这份详细的拆解能为你复现或开展类似研究铺平道路。本文还有配套的精品资源点击获取