Vision Transformer图像去雾：全局自注意力原理与PyTorch工程实践

简介：面向图像去雾与视觉Transformer研究者和开发者，提供基于Vision Transformer算法的完整Python实现，含源代码、预训练权重与详细使用说明。压缩包共340个文件，约156MB，其中204个Python脚本负责模型构建、训练参数配置与推理流程，16个YAML文件定义实验环境与超参数，12个CSV记录训练日志及评价指标，9个Jupyter Notebook提供交互式演示，另有GIF对比图与效果图辅助理解，方便直观比较算法性能。已有1446人浏览学习。资源附带了按不同数据集划分的预训练权重（My_best_model文件夹），可通过--pretrain_weights参数指定权重路径继续训练；训练所需的补丁大小（--train_ps）等核心参数集中在option.py中设置，便于复现实验。此外，项目介绍文档和详细使用说明覆盖了从环境配置到模型评估的完整链路，适合有一定深度学习基础、想深入理解ViT在图像去雾低层视觉任务中应用的开发者。

1. 为什么去雾任务越来越多人转向Vision Transformer而非纯卷积网络

监控视频、无人机航拍、自动驾驶感知，遇到雾天基本全靠后期算法救。过去几年，暗通道先验和CNN去雾模型是主流，但做过落地的人都有一个共同体验：模型在合成雾评测集上PSNR能到30dB以上，一到真实雾图就翻车。这个现象背后的原因很反直觉——CNN引以为傲的局部归纳偏置，恰恰在去雾任务上成了包袱。雾的影响是全局的，同一张图里近景清晰、远景发白，透射率分布跨越整个空间；Vision Transformer通过全局自注意力把每个图像块和全图所有块建立联系，天然适合拟合这种长程、浓度渐变的物理过程。这篇笔记要讲清楚：基于Vision Transformer的图像去雾算法怎么做，Python工程从数据准备、模型搭建到训练推理怎么跑通，以及那些不跑一遍根本发现不了的坑。

2. 从大气散射模型到全局自注意力：去雾任务的原理与选型理由

2.1 大气散射模型：去雾问题到底在解什么

去雾的物理基础是大气散射模型，写成公式是 I(x) = J(x)t(x) + A(1 - t(x))。I是相机收到的模糊图，J是要恢复的干净图，t是透射率（0到1之间），A是全局大气光。雾越浓，t越小，远端像素的贡献被A压得越厉害。算法要么估计t和A再逆推J，要么端到端直接回归J。问题在于，一张模糊图对应无穷多组(J, t, A)，这是个病态反问题，模型必须靠数据和先验把解空间压住，否则恢复结果会偏色或者丢失细节。

早年方法用先验约束t，比如暗通道先验假设局部小块里至少一个颜色通道接近0，最初的去雾效果确实惊艳，但遇到天空、白色物体、明亮灯光这些违反先验的区域就露馅。这类方法的输出常常带一圈黑边，或者把白色墙壁恢复成灰紫色。根本原因在于它用局部统计量估计全局物理量，天然缺少跨区域推理能力。CNN方法好一些，但受感受野限制，想要建模远处浓雾和近景清晰的连续分布，需要很深的网络不断膨胀感受野，训练难度和推理成本都上升，而且理论感受野不等于有效感受野——深层卷积真正依赖的区域往往远小于理论值，远距离像素之间的关联仍然很弱。

2.2 ViT的自注意力凭什么能拟合透射率分布

Vision Transformer把图像切成固定大小的patch，每个patch拉平成token，经过一层线性投影后加上位置编码，送进Transformer Encoder。核心是自注意力，每一层里每个token都会和全图所有token计算相关性，输出权重加权融合。这意味着第一层就能看到整张图的上下文，远处天空的信息可以直接影响近景地面块的估计，不需要靠层数堆叠。对去雾任务，这个性质非常关键：透射率图通常是全局平滑、局部突变的混合结构，全局注意力负责捕捉平滑渐变，前馈网络负责拟合局部突变，两类信息在一个模块里共同决策。

具体看Query、Key、Value的角色：某个patch的Query和其他所有patch的Key做点积，得到注意力权重，再去加权对应的Value。响应值越高的token，对当前patch的恢复结果影响越大。模型训练的过程，就是学会让远处浓度相似的patch互相增强、让边界处反差大的patch互相抑制。位置编码还保留了patch之间的空间关系，不会因为全局注意力丢掉图像结构。与分类任务不同，去雾不用CLS token做全局表征，而是保留每个token的输出，reshape回二维特征图后再逐像素回归，这样输出天然对齐输入位置，不需要额外的对齐模块。

2.3 与CNN、基于窗口注意力的Transformer变体对比：什么时候选ViT不至于踩坑

CNN在小雾、局部雾场景下已经够用，且推理快、显存低，但感受野是硬伤。基于窗口注意力的Transformer变体用窗口注意力把自注意力限制在局部窗口内，计算量降了，但跨窗口信息需要靠移位窗口来交换，实现复杂度和调参难度都上升。纯ViT在这个任务上的优势是趁手：它吃全局建模，patch_size和层数可以直接参考成熟分类配置，上手路径清晰；代价是长序列注意力显存占用高，训练收敛比CNN慢。

我的选择习惯是：白天均匀雾、数据量大，用ViT做主干端到端回归；夜间或非均匀雾，ViT结构保留，但要额外加物理模型先验分支。从复现一个开源Python工程的角度，ViT版本也比窗口注意力变体更容易读懂梯度，因为每个token都直达全图，backward路径没有窗口边界，调试时少一层玄学。

3. 环境搭建与数据预处理：先让模型跑起来再说

3.1 去雾数据集怎么组织：训练、验证、测试目录结构与格式

常见做法是准备三个目录：train存放成对的模糊图和清晰图；val用来每N个epoch算PSNR；test放真实雾图，只用清晰图做参考或不做参考。为解决合成雾和真实雾的domain gap，实际工程会把合成雾数据当作主训练集，再用少量真实雾图做微调。目录组织示例：

data/ ├── train/ │ ├── hazy/ # 模糊图，jpg 或 png │ └── clear/ # 对应的清晰图，文件名与 hazy 一一对应 ├── val/ │ ├── hazy/ │ └── clear/ └── test/ └── real_world/ # 真实雾图，无配对清晰图

文件逐一配对是这里最容易翻车的地方。我见过不止一次因为train和clear里文件名顺序不一致，模型学到的是「把上一张图变成下一张图的颜色」，PSNR看着高，换新图全糊。所以建议加载时用dict按文件名建立映射，而不是拿两个列表zip。此外，建议把所有图统一缩放，长边不超过1024，训练时用随机裁剪，不要整图送入ViT——整图输入会让patch数量激增，显存直接爆掉。

验证集计算PSNR时要保持输入尺寸一致，不能用训练时的随机裁剪，否则每次验证的裁剪位置不同，指标波动大。常见做法是固定中心裁剪到256×256，或者对整图做滑窗推理再拼回来。我一般会把val集按浓度分三组——轻雾、中雾、重雾，分别统计PSNR，避免一组数据拉高平均值掩盖重雾区域的失效。

3.2 预处理参数：归一化、patch_size和数据增广的取舍

ViT对输入有一个硬性要求：图的高和宽必须是patch_size的整数倍。所以核心参数先定patch_size。去雾任务里我一般不用ViT-B/16那个16，而是用patch_size=8，后面接一层浅卷积做下采样任务。这会让token数量变为16倍，显存上升，但对细节恢复明显更好。如果显存吃紧，可以用16先跑通，再降到8微调。归一化直接用视觉预训练里最通用的mean=[0.485, 0.462, 0.406]、std=[0.229, 0.224, 0.225]，虽然这是分类任务的统计量，但预训练权重就是按这个范围对齐的，迁移效果最好，自己重新算一套统计量反而会让预训练失效。

注意：模糊图和清晰图必须使用相同的mean/std。如果分别归一化，相当于把清晰图的亮度基准改掉了，模型永远学不到正确的颜色映射。

数据增广方面，随机水平翻转、随机裁剪到256×256或320×320、色彩抖动（亮度/对比度/饱和度各±0.2）就够。去雾任务不建议做随机擦除或强旋转，因为雾的物理分布会被破坏，模型会学歪。训练时把模糊图和清晰图用同一个随机种子做相同的裁剪与翻转，保证像素严格对齐。这里我用一个自定义Dataset来实现配对增广，而不是直接在transform里各做各的。

3.3 训练入口的骨架：最小可运行的启动脚本

把训练循环封装成一个train.py，入口部分大概是这个样子：

import torch from torch.utils.data import DataLoader from dataset import DehazeDataset from model import DehazeViT def main(): torch.manual_seed(0) device = 'cuda' if torch.cuda.is_available() else 'cpu' train_ds = DehazeDataset('data/train', patch_size=8, crop_size=256) train_dl = DataLoader(train_ds, batch_size=8, shuffle=True, num_workers=4, pin_memory=True, drop_last=True) model = DehazeViT(patch_size=8, embed_dim=192, depth=12, num_heads=3).to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4) for epoch in range(1, 301): for hazy, clear in train_dl: hazy = hazy.to(device) clear = clear.to(device) pred = model(hazy) loss = criterion(pred, clear) optimizer.zero_grad() loss.backward() optimizer.step()

这个骨架里，crop_size统一了输入分辨率，patch_size必须在数据集类和模型类里保持一致。batch_size=8是ViT-B/8配置下比较平衡的值，显存不够先减到4，同时开梯度累积。优化器AdamW配合weight_decay=1e-4是Transformer惯用组合，和CNN常用的SGD+momentum差别很大，直接套CNN的经验常常不收敛。

学习率调度器我推荐cosine退火，从1e-4降到1e-6，总训练轮数300到400。ViT训练慢主要体现在前期，前50个epoch可能看不出明显效果，不要因为验证集没涨就急着降低学习率，先确认loss在下降。训练日志里除了打印loss，建议每50步打印一次当前学习率和显存占用，方便定位训练停滞是因为学习率还是因为OOM。

4. 模型实现细节：从Patch Embedding到去雾输出头

4.1 核心结构：把ViT改造成回归模型的三个改动

去雾ViT不能直接拿来分类，要做三处改动。第一，输入头用卷积做Patch Embedding而不是单纯切片，重叠卷积能缓解patch边缘的不连续感。第二，位置编码因为训练裁剪尺寸固定而长度固定，推理不同尺寸时需要插值处理。第三，输出头把Encoder输出reshape回特征图，再接几个上采样卷积恢复原分辨率。简化实现如下：

import torch import torch.nn as nn class PatchEmbed(nn.Module): def __init__(self, in_ch=3, embed_dim=192, patch_size=8): super().__init__() # 重叠卷积：stride 取 patch_size 的一半，保留相邻 patch 重叠信息 self.proj = nn.Conv2d(in_ch, embed_dim, kernel_size=3, stride=patch_size // 2, padding=1) def forward(self, x): x = self.proj(x) # (B, C, H/4, W/4) 当 patch_size=8 时 x = x.flatten(2).transpose(1, 2) return x # (B, L, C) class DehazeHead(nn.Module): def __init__(self, embed_dim=192, out_ch=3): super().__init__() self.head = nn.Sequential( nn.Conv2d(embed_dim, embed_dim, 3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(embed_dim, out_ch, 3, padding=1), ) def forward(self, x, h, w): B, L, C = x.shape x = x.transpose(1, 2).reshape(B, C, h, w) return self.head(x)

Patch Embedding用卷积做的好处是保留局部纹理，stride取patch_size一半时相邻patch有重叠，输出特征图的高宽是输入的1/4而不是1/8，给输出头保留了更多空间细节。输出头最简可以只接两层3×3卷积直接回归RGB值，实测效果和带亚像素上采样的版本差不了太多，但建议用双线性插值把特征图放大到原始尺寸再卷积，边缘更干净。

Transformer Encoder的depth和num_heads不用刻意调大。depth=12、embed_dim=192、num_heads=3这个组合是显存和效果之间的平衡点。num_heads不是越多越好，每个head的维度通常建议保持64左右，192维拆3个head正好。qkv的bias建议保留，对回归任务有一点稳定作用，虽然分类任务里经常去掉。

4.2 损失函数怎么配：L1还是MSE，感知损失加不加

去雾训练的损失组合决定最终画质。纯MSE收敛快但图偏平滑，细节糊；纯L1能保留更多锐利边缘，但训练早期梯度震荡。常见做法是L1为主，加上SSIM损失和感知损失。感知损失不直接对VGG输出算L1，而是取预训练分类网络几个中间层特征分别算L1再加权，低层权重小、高层权重大，这样既约束颜色又约束结构。

def dehaze_loss(pred, target): l1 = torch.mean(torch.abs(pred - target)) ssim_loss = 1 - ssim(pred, target) perc = perceptual_loss(pred, target) # 多层特征 L1 加权 return 1.0 * l1 + 0.5 * ssim_loss + 0.1 * perc

权重不是拍脑袋定的。先全部跑1个epoch看每个分量的量级，把SSIM和感知损失的量级拉到和L1接近，再微调系数。比如L1在0.1量级，SSIM损失在0.05量级，感知损失如果只有0.01，那感知损失权重0.1就是合理的；如果感知损失初始就有0.1，权重需要降到0.01以下。如果发现训练后期PSNR不再涨，优先检查是不是L1权重太大把感知损失压没了。

4.3 训练循环与Checkpoint策略

训练循环需要留意三点：混合精度、EMA和checkpoint的保存策略。ViT训练很吃显存，自动混合精度能省将近一半显存，同时加速。EMA是对模型参数做滑动平均，推理时用EMA参数比用step后的参数更稳，去雾这类回归任务收益非常明显。Checkpoint不要只存model.state_dict()，还要存optimizer、epoch、EMA参数、当前学习率，否则恢复训练时优化器状态丢失，学习率调度不平滑。

scaler = torch.cuda.amp.GradScaler() ema_model = copy.deepcopy(model) ema_alpha = 0.999 for epoch in range(1, 301): model.train() for hazy, clear in train_dl: hazy, clear = hazy.to(device), clear.to(device) optimizer.zero_grad() with torch.cuda.amp.autocast(): pred = model(hazy) loss = dehaze_loss(pred, clear) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() with torch.no_grad(): for p, ema_p in zip(model.parameters(), ema_model.parameters()): ema_p.data.mul_(ema_alpha).add_(p.data, alpha=1 - ema_alpha) if epoch % 50 == 0: torch.save({ 'epoch': epoch, 'model': model.state_dict(), 'ema': ema_model.state_dict(), 'optimizer': optimizer.state_dict(), 'scaler': scaler.state_dict(), }, f'checkpoints/dehaze_vit_{epoch}.pt')

EMA固定alpha=0.999在小数据集够用，但严格做法是按step数修正：alpha = min(0.999, (step + 1) / (step + 10))，训练初期让EMA跟得快一点。保存间隔建议前50个epoch每10个epoch存一次，后面每50个存一次，这样早期翻车还能找回能用的中间权重。恢复训练时要把scaler也恢复进去，否则混合精度状态不一致，偶尔会出现loss突然跳到NaN又恢复正常的情况，看起来像玄学，其实只是GradScaler的state丢了。

5. 去雾训练与推理最常见的5个坑：现象、原因与解决

5.1 Loss停在0.1附近，PSNR卡在22-23dB不再上涨

现象：训练到约50个epoch，L1 loss在0.1左右徘徊，验证集PSNR长期停在22-23dB，怎么加epoch都没用。

原因：多半是全局学习率不合适，ViT对学习率比CNN敏感得多，1e-4配合AdamW如果配错了调度器，后期步长太大在局部震荡；另外输入图像范围没有对齐，如果图像是0-255而模型输出被sigmoid限制在0-1，loss永远降不到位。排查方法很简单：打印模型输出的像素值范围，看是否和target一致。

解决：先把数据统一到0-1范围，再用cosine退火调度器，初始lr降到3e-5左右重新训练。跑两三个epoch看验证PSNR是否有明显起色，再决定要不要继续。如果两三个epoch后PSNR还在22dB以下，检查数据配对的文件名是否一致，这个坑会让人白跑一周。

5.2 浓雾区域出现色块和伪影

现象：恢复出来的图远处浓雾区域有规则色块，边界和patch边界重合，像马赛克一样。

原因：Patch Embedding直接切patch，patch之间没有信息交换，位置编码每块一样，导致patch边缘不连续；输出头太浅，没法缝合。这个现象在patch_size=16时特别明显，patch越大色块越粗。

解决：把Patch Embedding改成重叠卷积，输出头加一层上采样卷积，或者在Transformer Encoder后面接一个3×3的卷积修边。调完再看浓雾区域的PSNR，一般能提升1-1.5dB。如果色块只出现在某些固定位置，还要检查位置编码是否在训练时被冻结或者被错误广播。

5.3 推理时换了图片尺寸直接崩溃

现象：训练时图是256×256，推理时传入512×384，报position embedding维度不匹配，或者shape错误。

原因：ViT固定长度位置编码只支持训练尺寸。这是ViT系列最经典的坑，任何用预训练ViT做下游任务的项目都会遇到。

解决：推理前把图padding到patch_size的整数倍，再插值位置编码。常见做法是：

ph = round(H / patch_size) * patch_size pw = round(W / patch_size) * patch_size x = F.interpolate(x, size=(ph, pw), mode='bilinear', align_corners=False) # 推理完成后，把输出裁剪回 (H, W)

更彻底的做法是在模型forward里对位置编码做双线性插值，这样任意尺寸都能跑，代价是每换一次尺寸多一次小计算。我倾向于后者，把逻辑内聚到模型里，这样导出ONNX时也能一起固化。

5.4 显存溢出，但batch_size明明设得很小

现象：batch_size=4、分辨率256×256还是OOM，而同样配置的CNN完全没压力。

原因：ViT自注意力的复杂度是O(n²)，n=token数量。patch_size=8时256×256有1024个token，每个token要和1024个token计算注意力，显存消耗远超CNN。此外混合精度没开也是常见助推因素，一个batch的反向传播会把所有中间注意力矩阵都保存在显存里。

解决：优先开启AMP梯度缩放；其次把patch_size从8改到16，token数降为原来的1/4；再不行就在Encoder层之间插入窗口注意力，或者用梯度累积把batch_size降到1-2。梯度累积在实现上要留意：loss除以累积步数后再backward，否则等效学习率会变大，模型训练不稳定。

5.5 夜间和强光源场景全面失效

现象：白天有雾图效果很好，换成夜间、路灯或车灯强光源场景，输出偏色严重，暗部死黑。

原因：大气散射模型假设大气光是全局均匀的，夜间有局域强光源，这个假设不成立；模型在合成雾数据上学到的分布里没有这类样本，属于分布外输入。

解决：如果只是白天场景落地，直接限制使用范围；如果非得做夜间，要在模型里加一个物理先验分支，把大气光A改成可学习的空间变化大气光图，并在训练集里加入夜间强光源合成样本。这个改动对工程复杂度是数量级的提升，我一般会先和需求方确认清楚：夜间是否真有硬性要求。很多时候做演示用的夜间去雾，最后用图像增强和色调映射就解决了，不必硬上物理模型。

6. 进阶：把去雾ViT导出成ONNX，做一套不依赖训练框架的推理脚本

模型训练完要落地，不能每次推理都开着Python训练环境。常见做法是导出ONNX。导出前先确认三点：模型处于eval模式；把位置编码插值逻辑固化到forward里，避免导出时动态shape导致算子不支持；输入输出都能用固定维度描述。用torch.onnx.export把模型导出为dehaze.onnx，opset_version选11以上，因为低版本对interpolate等算子的支持不稳定。

import numpy as np import onnxruntime as ort from PIL import Image sess = ort.InferenceSession('dehaze.onnx', providers=['CUDAExecutionProvider', 'CPUExecutionProvider']) def dehaze_infer(img_path, patch_size=8): img = Image.open(img_path).convert('RGB') H, W = img.size[1], img.size[0] ph = round(H / patch_size) * patch_size pw = round(W / patch_size) * patch_size img = img.resize((pw, ph), Image.BILINEAR) arr = np.asarray(img, dtype=np.float32) / 255.0 arr = (arr - [0.485, 0.462, 0.406]) / [0.229, 0.224, 0.225] x = arr.transpose(2, 0, 1)[None] # (1,3,H,W) out = sess.run(None, {sess.get_inputs()[0].name: x})[0] out = out[0].transpose(1, 2, 0) out = out * [0.229, 0.224, 0.225] + [0.485, 0.462, 0.406] out = np.clip(out, 0, 1) * 255.0 return out[:H, :W].astype(np.uint8)

这套脚本只依赖numpy、onnxruntime和PIL，生产环境不用装PyTorch。我在实际中吃过亏的地方是归一化细节：训练时用什么mean/std，推理必须完全一致，很多代码把这一步省了，导致输出整体偏色。还有一处要注意，ONNX Runtime的输入节点name在不同导出版本里会变，用sess.get_inputs()[0].name而不是写死一个字符串，能少踩一个坑。我现在的习惯是：训练脚本和推理脚本共用一个preprocess.py，所有归一化和padding逻辑都从同一处引用。验证模型是否导出正确，拿同一张图分别跑PyTorch和ONNX输出对比，最大像素差不超过1就说明转换无损。一套去雾工程真正可靠的验收方式不是看训练时PSNR，而是拿一批现场真实雾图跑一遍，让没参与建模的人看输出是否自然、有没有偏色，这一步比任何指标都更能发现问题。希望帮到你。

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