尧图网页设计 
简介这是一份基于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_size16得到(256/16)^2256个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_channels3, out_channels3, patch_size16, embed_dim768, depth12, num_heads8): super().__init__() self.stem nn.Sequential( nn.Conv2d(in_channels, 64, kernel_size3, stride2, padding1), # 256-128 nn.ReLU(inplaceTrue), nn.Conv2d(64, 128, kernel_size3, stride2, padding1), # 128-64 nn.ReLU(inplaceTrue), nn.Conv2d(128, embed_dim, kernel_size3, stride1, padding1), # 64-64 ) self.patch_embed nn.Conv2d(embed_dim, embed_dim, kernel_sizepatch_size, stridepatch_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_modelembed_dim, nheadnum_heads, dim_feedforwardembed_dim * 4, dropout0.1, activationgelu, batch_firstTrue) self.encoder nn.TransformerEncoder(encoder_layer, num_layersdepth) self.head nn.Sequential( nn.ConvTranspose2d(embed_dim, 128, kernel_size2, stride2), nn.ReLU(inplaceTrue), nn.Conv2d(128, 64, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.ConvTranspose2d(64, out_channels, kernel_size4, stride2, padding1), ) 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], dim1) 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_size16, embed_dim768, depth12, num_heads8) print(params:, sum(p.numel() for p in model.parameters() if p.requires_grad))这里的stem用三层卷积做下采样目的是在切patch前把低频信息和高频边缘都压到embedding里。patch_size16时64x64的特征图切出4x416个token相比直接对256x256切256个token注意力矩阵规模降了一个量级。embed_dim768、depth12、num_heads8对应ViT-Base配置训练时约8600万参数。如果显存吃紧可以改成embed_dim384、depth6参数少一半能跑但去雾边缘会稍微软一点。TransformerEncoderLayer里dim_feedforward我设成embed_dim*4这是ViT的默认配置。要注意batch_firstTrue这样tensor维度顺序是(B, seq, embed)否则你后续reshape会鬼打墙。3. 四种真实雾凇数据集NH-HAZE、NTIRE2019、I-HAZE与O-HAZE的选用逻辑3.1 四种数据集的差异数据集场景雾的类型数据量特点对训练的影响I-HAZE室内均匀雾相对较少patch间雾浓度差异小适合起步O-HAZE室外均匀雾场景深度差异中等覆盖远景退化泛化最关键NH-HAZE室内非均匀雾少难度最高无严格对齐NTIRE2019混合场景真实拍摄多测试泛化的金标准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写一份可跑的datasetimport 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_size256, is_trainTrue): 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), interpolationcv2.INTER_CUBIC) clean cv2.resize(clean, (256, 256), interpolationcv2.INTER_CUBIC) hazy torch.from_numpy(hazy.transpose(2, 0, 1)) clean torch.from_numpy(clean.transpose(2, 0, 1)) return hazy, clean逻辑说明读取雾图和对应的清晰图BGR转RGBfloat32归一化到[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_rootdata/I-HAZE/train/hazy, clean_rootdata/I-HAZE/train/GT, crop_size256, is_trainTrue), batch_size4, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastTrue ) val_loader DataLoader( DehazeDataset(hazy_rootdata/NTIRE2019/val/hazy, clean_rootdata/NTIRE2019/val/GT, crop_size256, is_trainFalse), batch_size1, shuffleFalse, num_workers2 )batch_size4对ViT-Base已经偏大后面会讲显存怎么控制。num_workers在Windows下建议设为0否则DataLoader会报worker错误。pin_memoryTrue能减少GPU拷贝但显存不足时可以先关掉。drop_lastTrue是为了防止最后一个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 python3.9 conda activate dehaze pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 pip install opencv-python scikit-image matplotlib tqdmcu118对应CUDA 11.8如果本机是CUDA 12.x就换成cu121。OpenCV用来读取图像scikit-image用来算PSNR和SSIMtqdm看训练进度。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(weightsVGG16_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_range1.0) # 由pytorch_msssim提供 return l1 0.05 * perceptual 0.2 * ssim_vall1_loss是像素级保真度的主力perceptual用VGG16前16层输出的MSE作用是让输出在语义结构上也接近清晰图ssim_loss直接优化结构相似性。perceptual权重0.05、ssim权重0.2是我从多组实验里拿出来的。如果只用L1输出边缘虽然锐利但色彩容易漂移只用L1perceptual整体色调对但局部细节容易糊尤其是草地和树叶这类高纹理区域。VGG的BatchNorm层在训练时也会更新running_mean/running_var所以冻结参数还不够还要调用model.eval()否则perceptual loss本身会越来越不稳定。我在这上面翻过车loss慢慢涨就是不降最后发现是VGG里的BN在动。4.3 优化器、学习率调度与EMAViT对学习率极其敏感稍有不慎就会loss炸到NaN。我用AdamWlr1e-4weight_decay0.05再配一个cosine decay调度和5轮warmupoptimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay0.05) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max300, eta_min1e-6)训练入口参数python train.py --model dehaze_vit --dataset I-HAZE --batch_size 4 --lr 1e-4 --epochs 300 --crop_size 256 --checkpoint ckpt/dehaze_vitmodel指定模型结构dataset指定四种数据集之一batch_size是每个GPU上的批大小lr是初始学习率epochs是总轮数。去雾属于低层视觉任务不像分类那样30轮就能收敛一般到150轮之后PSNR才会明显抬升前100轮看起来像没学别急着停。训练中的EMA指数移动平均对去雾很有效。我通常设decay0.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_size8batch_size8一张RTX 2080Ti 11G直接OOM。原因ViT的Self-Attention内存复杂度是O(n^2)n是token数。图像256x256patch16时token数是256patch8时token数是1024注意力矩阵直接多16倍显存从不到10G飙升到接近30G。batch_size8时8个序列一起算直接溢出。提示显存不够时优先减小batch_size保持patch_size16如果还需要更大感受野考虑window attention或Swin Transformer的局部注意力。解决batch_size4起步梯度累积两次等效batch8想要细节更好就把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(strictFalse)并做一层公共键过滤state_dict torch.load(vit_base.pth, map_locationcpu) 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, strictFalse)这段过滤逻辑是关键只在同名且同shape的键做拷贝避免head、pos_embed等维度不一致导致的意外。过滤完打印一下新state_dict的键数量如果只有零头说明前几个关键层的命名也匹配不上那就是骨干网络结构不完全一致得看代码里有没有做add_norm或pre_norm的差异。5.4 PSNR虚高数据范围没换算回[0, 255]现象第一次评估PSNR63dB把测试结果发群里被同行提醒数值不对劲。原因训练时把图像归一化到了[0,1]而PSNR的MSE是基于[0,255]的像素误差计算的。直接拿0~1的pred和target算MSE误差被缩小255^2倍PSNR虚高约48dB。解决算PSNR前先转回[0,255]def compute_psnr(pred, target, max_val255.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_ratiodata_range要传1.0而不是255。这是我踩过最蠢的坑但也是出现频率最高的坑很多号称模型无敌的截图其实是这一步没换算。5.5 梯度爆炸与NaN loss现象训练到第50轮loss突然变成nan再也没降下来。原因ViT对学习率极度敏感lr1e-3大概率炸另外fp16混合精度下激活值溢出也会nan。解决lr设为1e-4并带5轮warmup在backward之后加cliploss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.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), modebilinear).squeeze() return attn_map.cpu().numpy()hook捕获的是注意力模块的输出output[1]是注意力权重矩阵。我跑多个数据集后得到一个规律当模型真正学到去雾时注意力热力图会集中在浓雾与边缘交界区域如果热力图呈无规律的均匀噪声说明模型根本没在学去雾只是在做全局亮度调整。这个检查比PSNR早发现模型走偏。6.3 验证模型是真去雾还是假去雾高PSNR并不代表一切。有个我印象很深的教训有一次训练出来的模型在O-HAZE上PSNR 29dB但逐张看图发现暗部被整体提亮天空颜色过度饱和。PSNR不低是因为误差区域只占小部分但视觉观感非常差。后来我习惯每10轮就把一个batch的预测结果拼成四格图有雾图、GT、输出、注意力热力图并排看同时检查纹理细节、色彩漂移和注意力分布。从那以后我每次训练完必定先跑一遍可视化再算PSNR和SSIM——不再把单一指标当成衡量模型好坏的唯一标准。希望帮到你。本文还有配套的精品资源点击获取 基于Vision Transformer的图像去雾算法：从原理到实战
简介这是一份基于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_size16得到(256/16)^2256个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_channels3, out_channels3, patch_size16, embed_dim768, depth12, num_heads8): super().__init__() self.stem nn.Sequential( nn.Conv2d(in_channels, 64, kernel_size3, stride2, padding1), # 256-128 nn.ReLU(inplaceTrue), nn.Conv2d(64, 128, kernel_size3, stride2, padding1), # 128-64 nn.ReLU(inplaceTrue), nn.Conv2d(128, embed_dim, kernel_size3, stride1, padding1), # 64-64 ) self.patch_embed nn.Conv2d(embed_dim, embed_dim, kernel_sizepatch_size, stridepatch_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_modelembed_dim, nheadnum_heads, dim_feedforwardembed_dim * 4, dropout0.1, activationgelu, batch_firstTrue) self.encoder nn.TransformerEncoder(encoder_layer, num_layersdepth) self.head nn.Sequential( nn.ConvTranspose2d(embed_dim, 128, kernel_size2, stride2), nn.ReLU(inplaceTrue), nn.Conv2d(128, 64, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.ConvTranspose2d(64, out_channels, kernel_size4, stride2, padding1), ) 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], dim1) 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_size16, embed_dim768, depth12, num_heads8) print(params:, sum(p.numel() for p in model.parameters() if p.requires_grad))这里的stem用三层卷积做下采样目的是在切patch前把低频信息和高频边缘都压到embedding里。patch_size16时64x64的特征图切出4x416个token相比直接对256x256切256个token注意力矩阵规模降了一个量级。embed_dim768、depth12、num_heads8对应ViT-Base配置训练时约8600万参数。如果显存吃紧可以改成embed_dim384、depth6参数少一半能跑但去雾边缘会稍微软一点。TransformerEncoderLayer里dim_feedforward我设成embed_dim*4这是ViT的默认配置。要注意batch_firstTrue这样tensor维度顺序是(B, seq, embed)否则你后续reshape会鬼打墙。3. 四种真实雾凇数据集NH-HAZE、NTIRE2019、I-HAZE与O-HAZE的选用逻辑3.1 四种数据集的差异数据集场景雾的类型数据量特点对训练的影响I-HAZE室内均匀雾相对较少patch间雾浓度差异小适合起步O-HAZE室外均匀雾场景深度差异中等覆盖远景退化泛化最关键NH-HAZE室内非均匀雾少难度最高无严格对齐NTIRE2019混合场景真实拍摄多测试泛化的金标准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写一份可跑的datasetimport 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_size256, is_trainTrue): 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), interpolationcv2.INTER_CUBIC) clean cv2.resize(clean, (256, 256), interpolationcv2.INTER_CUBIC) hazy torch.from_numpy(hazy.transpose(2, 0, 1)) clean torch.from_numpy(clean.transpose(2, 0, 1)) return hazy, clean逻辑说明读取雾图和对应的清晰图BGR转RGBfloat32归一化到[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_rootdata/I-HAZE/train/hazy, clean_rootdata/I-HAZE/train/GT, crop_size256, is_trainTrue), batch_size4, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastTrue ) val_loader DataLoader( DehazeDataset(hazy_rootdata/NTIRE2019/val/hazy, clean_rootdata/NTIRE2019/val/GT, crop_size256, is_trainFalse), batch_size1, shuffleFalse, num_workers2 )batch_size4对ViT-Base已经偏大后面会讲显存怎么控制。num_workers在Windows下建议设为0否则DataLoader会报worker错误。pin_memoryTrue能减少GPU拷贝但显存不足时可以先关掉。drop_lastTrue是为了防止最后一个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 python3.9 conda activate dehaze pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 pip install opencv-python scikit-image matplotlib tqdmcu118对应CUDA 11.8如果本机是CUDA 12.x就换成cu121。OpenCV用来读取图像scikit-image用来算PSNR和SSIMtqdm看训练进度。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(weightsVGG16_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_range1.0) # 由pytorch_msssim提供 return l1 0.05 * perceptual 0.2 * ssim_vall1_loss是像素级保真度的主力perceptual用VGG16前16层输出的MSE作用是让输出在语义结构上也接近清晰图ssim_loss直接优化结构相似性。perceptual权重0.05、ssim权重0.2是我从多组实验里拿出来的。如果只用L1输出边缘虽然锐利但色彩容易漂移只用L1perceptual整体色调对但局部细节容易糊尤其是草地和树叶这类高纹理区域。VGG的BatchNorm层在训练时也会更新running_mean/running_var所以冻结参数还不够还要调用model.eval()否则perceptual loss本身会越来越不稳定。我在这上面翻过车loss慢慢涨就是不降最后发现是VGG里的BN在动。4.3 优化器、学习率调度与EMAViT对学习率极其敏感稍有不慎就会loss炸到NaN。我用AdamWlr1e-4weight_decay0.05再配一个cosine decay调度和5轮warmupoptimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay0.05) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max300, eta_min1e-6)训练入口参数python train.py --model dehaze_vit --dataset I-HAZE --batch_size 4 --lr 1e-4 --epochs 300 --crop_size 256 --checkpoint ckpt/dehaze_vitmodel指定模型结构dataset指定四种数据集之一batch_size是每个GPU上的批大小lr是初始学习率epochs是总轮数。去雾属于低层视觉任务不像分类那样30轮就能收敛一般到150轮之后PSNR才会明显抬升前100轮看起来像没学别急着停。训练中的EMA指数移动平均对去雾很有效。我通常设decay0.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_size8batch_size8一张RTX 2080Ti 11G直接OOM。原因ViT的Self-Attention内存复杂度是O(n^2)n是token数。图像256x256patch16时token数是256patch8时token数是1024注意力矩阵直接多16倍显存从不到10G飙升到接近30G。batch_size8时8个序列一起算直接溢出。提示显存不够时优先减小batch_size保持patch_size16如果还需要更大感受野考虑window attention或Swin Transformer的局部注意力。解决batch_size4起步梯度累积两次等效batch8想要细节更好就把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(strictFalse)并做一层公共键过滤state_dict torch.load(vit_base.pth, map_locationcpu) 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, strictFalse)这段过滤逻辑是关键只在同名且同shape的键做拷贝避免head、pos_embed等维度不一致导致的意外。过滤完打印一下新state_dict的键数量如果只有零头说明前几个关键层的命名也匹配不上那就是骨干网络结构不完全一致得看代码里有没有做add_norm或pre_norm的差异。5.4 PSNR虚高数据范围没换算回[0, 255]现象第一次评估PSNR63dB把测试结果发群里被同行提醒数值不对劲。原因训练时把图像归一化到了[0,1]而PSNR的MSE是基于[0,255]的像素误差计算的。直接拿0~1的pred和target算MSE误差被缩小255^2倍PSNR虚高约48dB。解决算PSNR前先转回[0,255]def compute_psnr(pred, target, max_val255.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_ratiodata_range要传1.0而不是255。这是我踩过最蠢的坑但也是出现频率最高的坑很多号称模型无敌的截图其实是这一步没换算。5.5 梯度爆炸与NaN loss现象训练到第50轮loss突然变成nan再也没降下来。原因ViT对学习率极度敏感lr1e-3大概率炸另外fp16混合精度下激活值溢出也会nan。解决lr设为1e-4并带5轮warmup在backward之后加cliploss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.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), modebilinear).squeeze() return attn_map.cpu().numpy()hook捕获的是注意力模块的输出output[1]是注意力权重矩阵。我跑多个数据集后得到一个规律当模型真正学到去雾时注意力热力图会集中在浓雾与边缘交界区域如果热力图呈无规律的均匀噪声说明模型根本没在学去雾只是在做全局亮度调整。这个检查比PSNR早发现模型走偏。6.3 验证模型是真去雾还是假去雾高PSNR并不代表一切。有个我印象很深的教训有一次训练出来的模型在O-HAZE上PSNR 29dB但逐张看图发现暗部被整体提亮天空颜色过度饱和。PSNR不低是因为误差区域只占小部分但视觉观感非常差。后来我习惯每10轮就把一个batch的预测结果拼成四格图有雾图、GT、输出、注意力热力图并排看同时检查纹理细节、色彩漂移和注意力分布。从那以后我每次训练完必定先跑一遍可视化再算PSNR和SSIM——不再把单一指标当成衡量模型好坏的唯一标准。希望帮到你。本文还有配套的精品资源点击获取