简介：这是一份基于Vision Transformer架构的图像去雾算法研究资源包，面向深度学习、计算机视觉方向的研究者与工程师，尤其适合需要真实雾霾数据集进行模型训练和效果验证的读者。压缩包共342个文件，约156MB，核心包含大量Python模型脚本、YAML训练/验证配置文件、CSV实验数据、可视化图片与GIF动图，以及便于快速上手的Jupyter Notebook和说明文档。资源覆盖NH-HAZE、NTIRE2019、I-HAZE、O-HAZE四种真实雾霾数据集上的完整训练与测试流程，提供了损失景观、指标对比等量化结果，可帮助理解Vision Transformer在图像去雾任务中的收敛特性与性能表现，也可作为论文复现、课程设计或算法对比研究的参考资料。目前已有125人学习下载。

1. 基于Vision Transformer的去雾算法：这套资源能复现到什么程度

拿到一张雾天图像，用暗通道先验去雾，远处建筑边缘总有一坨灰白色残雾；换成基于Vision Transformer的去雾模型后，同样的场景能还原出清晰轮廓。这不是调参玄学，而是Transformer的全局注意力机制在全局雾分布建模上，天然比CNN的局部感受野更占优势。这份资源把NH-HAZE、NTIRE2019、I-HAZE、O-HAZE四种真实雾霾数据集、基于ViT的去雾模型训练与测试流程完整打包在一起，模型结构、训练脚本、测试脚本都齐了。适合两类人：一是做图像去雾毕设或论文复现的研究生，二是工程上想换掉CNN去雾方案、又不想从零凑数据和调参的开发者。接下来我把模型选型、数据预处理、训练参数、五个最容易翻车的坑按顺序逐一拆开。

2. 为什么是ViT而不是更深CNN：全局注意力如何覆盖雾的真实分布

2.1 雾的物理模型与CNN感受野的错位

去雾任务背后是大气散射模型：

I(x) = J(x) * t(x) + A * (1 - t(x))

其中I(x)是有雾图像，J(x)是待恢复的清晰场景辐射，t(x)是传输图，A是大气光。雾的退化与场景深度强相关：远处的像素传输率低，近处的高。这个退化在空间分布上不是局部的，一片拉伸的雾带可能从图像左下角贯穿到右上角。CNN的卷积核默认只能看到局部邻域，即便用膨胀卷积把感受野撑到几十个像素，信息也要逐层传递才能覆盖全图，浅层丢失的边缘细节在中深层很难找回来。

这是去雾和分类最底层的差异：分类只需要识别"这是什么"，去雾需要恢复"每个像素长什么样"。所以CNN堆深度的增益在去雾上不如在分类上明显。我自己跑过ResNet-50作为backbone的变体，在O-HAZE测试集上把深度从50层加到101层，PSNR只涨了0.2dB，但单卡训练时长几乎翻倍，这个收益比让我果断转向Transformer结构。

2.2 ViT的patch embedding与自注意力机制

ViT的核心操作是把图像拆成固定大小的patch序列，每个patch经过线性变换转成embedding，再送入Transformer block。拿ViT-Base来说，输入256x256的图像，patch_size=16，得到(256/16)^2=256个patch token，加上一个CLS token就是257个token。每个token长度为768维。

Transformer block里最关键的是Multi-Head Self-Attention：每个token产生Q、K、V三个向量，通过(QK^T)/sqrt(d)算token与token之间的相似度，再用softmax归一化得到注意力权重，最后乘V得到加权的特征。这意味着一层注意力就能让处于图像左上角的patch和右下角的patch直接交互，雾在空间上的长程关联不需要像CNN那样靠堆层数慢慢传递。

在去雾场景下，这个特性的价值是：雾通常是全场域的，远近退化程度差异很大，但相邻patch的雾浓度可能突然变化，比如一片云影。ViT的多头注意力可以让不同head分别关注浓雾区、过渡区、细节区，这是CNN需要复杂模块设计才能勉强接近的能力。

2.3 去雾ViT的实际选型：patch与编码器深度怎么配

项目里常见的是ViT-Base/16为主干，配合一个浅层CNN作为前馈特征提取。因为纯ViT在像素级任务上容易丢失高频细节，直接回归清晰图会出现模糊边缘。我一般会在stem部分加三层卷积，把图像从256x256先降到64x64的特征图，再切patch，这样patch内部的局部纹理已经编码在embedding里。

以下是一个典型的模型骨架：

import torch import torch.nn as nn class DehazeViT(nn.Module): def __init__(self, in_channels=3, out_channels=3, patch_size=16, embed_dim=768, depth=12, num_heads=8): super().__init__() self.stem = nn.Sequential( nn.Conv2d(in_channels, 64, kernel_size=3, stride=2, padding=1), # 256->128 nn.ReLU(inplace=True), nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1), # 128->64 nn.ReLU(inplace=True), nn.Conv2d(128, embed_dim, kernel_size=3, stride=1, padding=1), # 64->64 ) self.patch_embed = nn.Conv2d(embed_dim, embed_dim, kernel_size=patch_size, stride=patch_size) self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed = nn.Parameter(torch.zeros(1, (64 // patch_size) ** 2 + 1, embed_dim)) encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, dim_feedforward=embed_dim * 4, dropout=0.1, activation='gelu', batch_first=True) self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=depth) self.head = nn.Sequential( nn.ConvTranspose2d(embed_dim, 128, kernel_size=2, stride=2), nn.ReLU(inplace=True), nn.Conv2d(128, 64, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.ConvTranspose2d(64, out_channels, kernel_size=4, stride=2, padding=1), ) def forward(self, x): f = self.stem(x) # [B, 768, 64, 64] B, C, H, W = f.shape tokens = self.patch_embed(f).flatten(2).transpose(1, 2) # [B, (H/ps)^2, C] cls_tokens = self.cls_token.expand(B, -1, -1) tokens = torch.cat([cls_tokens, tokens], dim=1) tokens = tokens + self.pos_embed features = self.encoder(tokens) features = features[:, 1:, :].transpose(1, 2) # 去掉CLS features = features.reshape(B, C, H // patch_size, W // patch_size) return self.head(features) model = DehazeViT(patch_size=16, embed_dim=768, depth=12, num_heads=8) print("params:", sum(p.numel() for p in model.parameters() if p.requires_grad))

这里的stem用三层卷积做下采样，目的是在切patch前把低频信息和高频边缘都压到embedding里。patch_size=16时，64x64的特征图切出4x4=16个token，相比直接对256x256切256个token，注意力矩阵规模降了一个量级。embed_dim=768、depth=12、num_heads=8对应ViT-Base配置，训练时约8600万参数。如果显存吃紧，可以改成embed_dim=384、depth=6，参数少一半，能跑但去雾边缘会稍微软一点。

TransformerEncoderLayer里dim_feedforward我设成embed_dim*4，这是ViT的默认配置。要注意batch_first=True，这样tensor维度顺序是(B, seq, embed)，否则你后续reshape会鬼打墙。

3. 四种真实雾凇数据集：NH-HAZE、NTIRE2019、I-HAZE与O-HAZE的选用逻辑

3.1 四种数据集的差异

I-HAZE是室内均匀雾，所有区域雾浓度几乎一致，模型容易学到"整体提亮+对比度拉伸"这种偷懒解法。O-HAZE是室外均匀雾，但场景深度范围大，近处清晰远处模糊，模型必须学会按深度差异化去雾，否则远景边缘一定糊。NH-HAZE是非均匀雾霾，同一张图里半边浓雾半边薄雾，对模型挑战最大，它的浓雾区和薄雾区交界处是最能看出模型真本事的区域。NTIRE2019挑战赛数据集包含多种室内外真实拍摄场景，我一般把它留出来当验证集，避免训练时不小心泄露。

3.2 数据预处理：为什么要统一尺寸和归一化

这四种数据集原始分辨率不同。I-HAZE大约600x400附近，O-HAZE是640x480左右，NH-HAZE分辨率也不一致。ViT有位置编码，输入尺寸固定，所以在Dataset里要把训练样本统一到256x256。这里有一个关键选择：训练用随机裁剪而不是直接resize。

直接resize的问题在于，雾图像的大片背景经压缩后会出现蚊音噪声，评估PSNR时也会被这个小噪声干扰。随机裁剪到256x256，再配合水平翻转和颜色抖动，相当于把单张图的训练样本数量增加，ViT的泛化会更好。

归一化也容易踩坑。PyTorch预训练模型常用的ImageNet均值和标准差是mean=[0.485,0.456,0.406]、std=[0.229,0.224,0.225]，但去雾模型通常不需要完全对齐ImageNet分布，用简单归一化到[0,1]也可以。我的习惯是先把图像转到float32、除以255、再减均值除标准差。如果预训练encoder用的是ImageNet权重，那么前层统计分布是对齐的，加载预训练才有意义。

3.3 Dataset与DataLoader落地写法

用PyTorch写一份可跑的dataset：

import os import cv2 import numpy as np import torch from torch.utils.data import Dataset class DehazeDataset(Dataset): def __init__(self, hazy_root, clean_root, crop_size=256, is_train=True): self.hazy_paths = sorted(os.listdir(hazy_root)) self.clean_root = clean_root self.crop_size = crop_size self.is_train = is_train def __len__(self): return len(self.hazy_paths) def __getitem__(self, idx): hname = self.hazy_paths[idx] hazy = cv2.imread(os.path.join(self.hazy_root, hname)) clean = cv2.imread(os.path.join(self.clean_root, hname.replace('hazy', 'GT').replace('.jpg', '.png'))) hazy = cv2.cvtColor(hazy, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 clean = cv2.cvtColor(clean, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 if self.is_train: h, w = hazy.shape[:2] top = np.random.randint(0, h - self.crop_size + 1) left = np.random.randint(0, w - self.crop_size + 1) hazy = hazy[top:top + self.crop_size, left:left + self.crop_size] clean = clean[top:top + self.crop_size, left:left + self.crop_size] if np.random.rand() < 0.5: hazy = hazy[:, ::-1, :].copy() clean = clean[:, ::-1, :].copy() else: hazy = cv2.resize(hazy, (256, 256), interpolation=cv2.INTER_CUBIC) clean = cv2.resize(clean, (256, 256), interpolation=cv2.INTER_CUBIC) hazy = torch.from_numpy(hazy.transpose(2, 0, 1)) clean = torch.from_numpy(clean.transpose(2, 0, 1)) return hazy, clean

逻辑说明：读取雾图和对应的清晰图，BGR转RGB，float32归一化到[0,1]。训练时随机裁剪并做水平翻转，测试时统一resize到256x256；返回的tensor形状是(3, H, W)。参数说明：crop_size控制裁剪尺寸，越大ViT看到的上下文越全，但显存压力也大；hazy_paths和clean_paths的命名规则我写的是"hazy→GT"的映射，实际数据集命名千奇百怪，建议先打印两条路径确认配对，再跑DataLoader。

DataLoader配置如下，同样有讲究：

from torch.utils.data import DataLoader train_loader = DataLoader( DehazeDataset(hazy_root='data/I-HAZE/train/hazy', clean_root='data/I-HAZE/train/GT', crop_size=256, is_train=True), batch_size=4, shuffle=True, num_workers=4, pin_memory=True, drop_last=True ) val_loader = DataLoader( DehazeDataset(hazy_root='data/NTIRE2019/val/hazy', clean_root='data/NTIRE2019/val/GT', crop_size=256, is_train=False), batch_size=1, shuffle=False, num_workers=2 )

batch_size=4对ViT-Base已经偏大，后面会讲显存怎么控制。num_workers在Windows下建议设为0，否则DataLoader会报worker错误。pin_memory=True能减少GPU拷贝，但显存不足时可以先关掉。drop_last=True是为了防止最后一个batch的shape与其他batch不一致导致BatchNorm或注意力mask报错。

4. 训练与测试落地：环境配置、损失函数与脚本拆解

4.1 环境配置：PyTorch与CUDA的匹配

我建议Python 3.9 + PyTorch 2.0及以上，CUDA 11.8或12.1。用conda干净：

conda create -n dehaze python=3.9 conda activate dehaze pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 pip install opencv-python scikit-image matplotlib tqdm

cu118对应CUDA 11.8，如果本机是CUDA 12.x就换成cu121。OpenCV用来读取图像，scikit-image用来算PSNR和SSIM，tqdm看训练进度。ViT去雾没有比这些更特殊的依赖，最常见的深度学习环境配置坑主要来自torch版本与CUDA版本不匹配。装完后验证一下：

python -c "import torch; print(torch.cuda.is_available(), torch.__version__)"

返回True再继续，万一返回False，先查nvidia-smi的驱动版本，再决定装哪个cu版本。不要在驱动上硬刚，纯属浪费时间。

4.2 损失函数组合

去雾任务的损失函数，我最终用的是L1 Loss + Perceptual Loss + SSIM Loss的组合：

import torch.nn as nn import torch.nn.functional as F from torchvision.models import vgg16 class DehazeLoss(nn.Module): def __init__(self): super().__init__() vgg = vgg16(weights=VGG16_Weights.IMAGENET1K_V1).features[:16].eval() for p in vgg.parameters(): p.requires_grad = False self.vgg = vgg def forward(self, pred, target): l1 = F.l1_loss(pred, target) vgg_pred = self.vgg(pred) vgg_target = self.vgg(target) perceptual = F.mse_loss(vgg_pred, vgg_target) ssim_val = 1 - ssim(pred, target, data_range=1.0) # 由pytorch_msssim提供 return l1 + 0.05 * perceptual + 0.2 * ssim_val

l1_loss是像素级保真度的主力，perceptual用VGG16前16层输出的MSE，作用是让输出在语义结构上也接近清晰图，ssim_loss直接优化结构相似性。perceptual权重0.05、ssim权重0.2是我从多组实验里拿出来的。如果只用L1，输出边缘虽然锐利但色彩容易漂移；只用L1+perceptual，整体色调对，但局部细节容易糊，尤其是草地和树叶这类高纹理区域。

VGG的BatchNorm层在训练时也会更新running_mean/running_var，所以冻结参数还不够，还要调用model.eval()，否则perceptual loss本身会越来越不稳定。我在这上面翻过车，loss慢慢涨就是不降，最后发现是VGG里的BN在动。

4.3 优化器、学习率调度与EMA

ViT对学习率极其敏感，稍有不慎就会loss炸到NaN。我用AdamW，lr=1e-4，weight_decay=0.05，再配一个cosine decay调度和5轮warmup：

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.05) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=300, eta_min=1e-6)

训练入口参数：

python train.py --model dehaze_vit --dataset I-HAZE --batch_size 4 --lr 1e-4 --epochs 300 --crop_size 256 --checkpoint ckpt/dehaze_vit

model指定模型结构，dataset指定四种数据集之一，batch_size是每个GPU上的批大小，lr是初始学习率，epochs是总轮数。去雾属于低层视觉任务，不像分类那样30轮就能收敛，一般到150轮之后PSNR才会明显抬升，前100轮看起来像没学，别急着停。

训练中的EMA（指数移动平均）对去雾很有效。我通常设decay=0.999，每20个iteration把模型参数做一次平滑，验证阶段用EMA权重而不是原始权重评估。EMA相当于把上百个checkpoint的权重做了平均，能明显压低验证集上PSNR的抖动，让曲线更平滑。

4.4 测试与推理

验证函数长这样：

def evaluate(model, loader, device): model.eval() psnr_list, ssim_list = [], [] with torch.no_grad(): for hazy, clean in loader: hazy, clean = hazy.to(device), clean.to(device) pred = model(hazy) psnr_list.append(compute_psnr(pred, clean)) ssim_list.append(compute_ssim(pred, clean)) return float(np.mean(psnr_list)), float(np.mean(ssim_list))

逻辑说明：在no_grad下跑完整验证集，对每个batch算PSNR/SSIM再取平均。注意模型要切成eval模式，因为dropout和BatchNorm在训练/推理下行为不同。val_loader的batch_size建议是1，如果输入分辨率与位置编码不一致，ViT会直接报错。

保存模型我习惯分两份：

torch.save({ 'model': model.state_dict(), 'ema': ema.state_dict(), }, 'ckpt/last.pt')

EMA权重单独存一份，因为走EMA之后的权重比raw权重在去雾任务上PSNR高0.1~0.3dB，测试推理时优先加载ema。加载后务必调用model.eval()，输入要做同样的归一化，输出要反归一化再转存成图像。

5. 避坑指南：从显存溢出到PSNR虚高的五个翻车现场

5.1 显存溢出：patch_size与batch_size的双重压力

现象：patch_size=8，batch_size=8，一张RTX 2080Ti 11G直接OOM。

原因：ViT的Self-Attention内存复杂度是O(n^2)，n是token数。图像256x256，patch=16时token数是256；patch=8时token数是1024，注意力矩阵直接多16倍，显存从不到10G飙升到接近30G。batch_size=8时，8个序列一起算，直接溢出。

提示：显存不够时优先减小batch_size，保持patch_size=16；如果还需要更大感受野，考虑window attention或Swin Transformer的局部注意力。

解决：batch_size=4起步，梯度累积两次等效batch=8；想要细节更好就把patch_size降到8，此时batch_size减到2，同时开混合精度fp16。我现在做ViT去雾，显存不够第一反应永远是开fp16，而不是降batch，因为fp16对训练精度的影响在去雾这种回归任务上几乎不可见。

5.2 NH-HAZE非均匀雾导致验证指标横跳

现象：训练损失一直降，但NH-HAZE上的PSNR不涨反跌。

原因：NH-HAZE是非均匀雾，雾浓度在空间上变化剧烈，而且部分雾图与GT存在对齐误差。模型在均匀雾数据上学到的"按全局浓度一次除雾"的映射，在非均匀雾上完全失效。

解决：把NH-HAZE放在测试集而不是训练集，先让模型在I-HAZE上训练，O-HAZE上微调，最后在NH-HAZE上评估。如果你的任务就是非均匀去雾，那在训练NH-HAZE时要加大crop_size，并降低perceptual_loss权重，让模型更专注局部结构，而不是被VGG特征拉着去对齐全局语义。

5.3 预训练权重加载：key mismatch与维度错位

现象：torch.load后报RuntimeError: Error(s) in loading state_dict for DehazeViT: Missing key(s)。

原因：ViT-Base的ImageNet权重含一个1000类的分类头，而我的DehazeViT头是一个转置卷积头，键名完全不匹配；另一个可能是position embedding的shape不同，ImageNet是224x224输入、pos_embed是(197, 768)，而我改成了256x256输入、pos_embed是(17, 768)。

解决：load_state_dict(strict=False)，并做一层公共键过滤：

state_dict = torch.load('vit_base.pth', map_location='cpu') new_state = {k: v for k, v in state_dict.items() if k in model.state_dict() and v.shape == model.state_dict()[k].shape} model.load_state_dict(new_state, strict=False)

这段过滤逻辑是关键：只在同名且同shape的键做拷贝，避免head、pos_embed等维度不一致导致的意外。过滤完打印一下新state_dict的键数量，如果只有零头，说明前几个关键层的命名也匹配不上，那就是骨干网络结构不完全一致，得看代码里有没有做add_norm或pre_norm的差异。

5.4 PSNR虚高：数据范围没换算回[0, 255]

现象：第一次评估PSNR=63dB，把测试结果发群里，被同行提醒数值不对劲。

原因：训练时把图像归一化到了[0,1]，而PSNR的MSE是基于[0,255]的像素误差计算的。直接拿0~1的pred和target算MSE，误差被缩小255^2倍，PSNR虚高约48dB。

解决：算PSNR前先转回[0,255]：

def compute_psnr(pred, target, max_val=255.0): pred = pred.clamp(0, 1) * max_val target = target.clamp(0, 1) * max_val mse = F.mse_loss(pred, target) return 10 * torch.log10(max_val ** 2 / mse)

如果你用skimage的peak_signal_noise_ratio，data_range要传1.0而不是255。这是我踩过最蠢的坑，但也是出现频率最高的坑，很多号称"模型无敌"的截图其实是这一步没换算。

5.5 梯度爆炸与NaN loss

现象：训练到第50轮，loss突然变成nan，再也没降下来。

原因：ViT对学习率极度敏感，lr=1e-3大概率炸；另外fp16混合精度下激活值溢出也会nan。

解决：lr设为1e-4并带5轮warmup；在backward之后加clip：

loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step()

如果还想用fp16加速，给AMP加一个GradScaler，并每个step调用scaler.update()。梯度裁剪对ViT几乎是必选项，因为自注意力的梯度方差比较大，不小心就冲爆。

6. 评估与进阶：把注意力图当调试工具，比PSNR更早发现问题

6.1 严谨计算PSNR与SSIM

如果你自己写评估，别直接用库的默认参数。skimage的peak_signal_noise_ratio要求data_range参数，传入1.0代表图像范围是0~1，传入255代表0~255。很多开源代码把归一化后的图像直接丢进去，忘了传data_range，导致PSNR虚高40多dB。我的习惯是统一写一个接口，内部先clamp到[0,1]，再手动转回[0,255]，用最朴素的MSE公式算，避免库版本差异。SSIM同理，win_size默认是7，输入图像尺寸小于7时库会直接报错。还有通道问题：PSNR按像素计算所以RGB三通道无所谓，但SSIM应该逐通道算再取平均，scikit-image默认会转灰度，两种算出来的分数差不少。

6.2 注意力可视化

ViT最大的优势是自注意力权重可以直接可视化，不像CNN那样是个黑匣子。把最后一层Transformer Encoder的注意力权重取出来，将CLS token对其他token的attention向量reshape回patch网格，再上采样回原图尺寸，就是一张热力图：

def visualize_attention(model, hazy, device): model.eval() features = {} def hook_fn(module, input, output): # nn.MultiheadAttention返回(输出, 注意力权重) features['attn'] = output[1].detach() model.encoder.layers[-1].self_attn.register_forward_hook(hook_fn) with torch.no_grad(): pred = model(hazy.unsqueeze(0).to(device)) attn = features['attn'] # (1, heads, seq_len, seq_len) attn_map = attn[0, 0, 0, 1:] # 第一个head里CLS对所有patch的注意力 side = int((attn_map.size(0)) ** 0.5) attn_map = attn_map.reshape(side, side) attn_map = F.interpolate(attn_map.unsqueeze(0).unsqueeze(0), size=(256, 256), mode='bilinear').squeeze() return attn_map.cpu().numpy()

hook捕获的是注意力模块的输出，output[1]是注意力权重矩阵。我跑多个数据集后得到一个规律：当模型真正学到去雾时，注意力热力图会集中在浓雾与边缘交界区域；如果热力图呈无规律的均匀噪声，说明模型根本没在学去雾，只是在做全局亮度调整。这个检查比PSNR早发现模型走偏。

6.3 验证模型是真去雾还是假去雾

高PSNR并不代表一切。有个我印象很深的教训：有一次训练出来的模型在O-HAZE上PSNR 29dB，但逐张看图发现暗部被整体提亮，天空颜色过度饱和。PSNR不低是因为误差区域只占小部分，但视觉观感非常差。后来我习惯每10轮就把一个batch的预测结果拼成四格图：有雾图、GT、输出、注意力热力图并排看，同时检查纹理细节、色彩漂移和注意力分布。

从那以后，我每次训练完，必定先跑一遍可视化，再算PSNR和SSIM——不再把单一指标当成衡量模型好坏的唯一标准。希望帮到你。

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