ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

VisionTransformer图像去雾实战:从原理到PyTorch源码复现

VisionTransformer图像去雾实战:从原理到PyTorch源码复现 简介这份资源是面向计算机相关专业学生与深度学习实践者的图像去雾项目基于Vision Transformer架构实现去雾算法适合作为毕业设计、课程设计或期末大作业的参考方案难度适中兼顾算法理解与工程落地。压缩包共338个文件约156.35MB以204个Python源码为核心辅以yaml配置、csv实验记录、ipynb笔记、png与gif效果图、md说明文档及pptx汇报材料覆盖训练、推理、评估与展示全流程。项目经导师指导并获评审98分源码均本地编译调试可运行已有102人学习。读者可从中获得完整的ViT去雾网络实现、数据处理与训练脚本、实验对比记录、可视化结果及文档说明便于快速复现、二次开发与论文写作参考。1. 从「雾里看花」到「一键清晰」VisionTransformer 图像去雾到底在做什么如果你手头有一批雾天拍回来的图客户催着要「能看清」的版本传统暗通道先验调参调到怀疑人生那基于 VisionTransformer 的图像去雾算法就是值得你花一个下午跑通的方向。它做的事很直接输入一张被雾霾降质的 RGB 图输出一张对比度恢复、色彩不偏、细节不糊的清晰图。和早期 DCP、AOD-Net 那类方案比ViT 去雾的核心优势在于全局建模——雾的分布往往不是局部均匀的远处浓、近处淡卷积核感受野再大也容易被局部纹理带偏而自注意力能一次性把整张图的透射率关系拉通。这套 python 源码加文档说明的形态适合两类人一是做毕设或课程设计、需要一份能跑通且有方法论述的完整工程二是已经写过 CNN 去雾、想对比 Transformer 结构到底强在哪的算法工程师。下面我按「先立住原理、再动手复现、最后讲坑」的顺序把这条链路拆开讲清楚。2. VisionTransformer 去雾的网络结构编码器、解码器与透射率估计怎么摆2.1 为什么去雾任务适合用 ViT 而不是纯 CNN去雾本质是一个图像到图像的回归问题但它和超分、去噪有个关键差别雾的浓度和场景深度强相关而深度是全局量。一张图里近处地面可能几乎无雾远处山体却白茫茫一片卷积网络要靠堆叠很多层才能把这种长距离依赖传过去中间还容易被局部高频纹理干扰。ViT 把图像切成 patch 后做自注意力任意两个 patch 之间直接建立联系透射率图 t(x) 和大气光 A 的估计就能在全局一致的前提下做出来。常见做法是采用「编码器-解码器 跳连」的混合结构编码器用若干 Transformer Block 提取多尺度特征解码器逐步上采样回原分辨率跳连把浅层的边缘细节补回去。纯 ViT 直接做密集预测会有 patch 边界伪影所以实际工程里我一般会在解码末端加两三个卷积层做平滑这一步对最终 PSNR 影响不小。2.2 一个可复现的编码器-解码器骨架下面这段是网络主体的最小可跑版本用 PyTorch 写patch 大小设 8嵌入维度 256编码器 6 个 Block解码器对称。你可以直接抄进自己的工程改。import torch import torch.nn as nn from einops import rearrange class PatchEmbed(nn.Module): 把 HxWx3 的图切成 patch 并线性映射到 embed_dim def __init__(self, img_size256, patch_size8, in_chans3, embed_dim256): super().__init__() self.patch_size patch_size # 用卷积实现切块映射比 unfold 更快 self.proj nn.Conv2d(in_chans, embed_dim, kernel_sizepatch_size, stridepatch_size) def forward(self, x): x self.proj(x) # B, C, H/p, W/p x rearrange(x, b c h w - b (h w) c) return x class TransformerBlock(nn.Module): 标准 ViT BlockLayerNorm - MHSA - 残差 - MLP - 残差 def __init__(self, dim, num_heads8, mlp_ratio4.0): super().__init__() self.norm1 nn.LayerNorm(dim) self.attn nn.MultiheadAttention(dim, num_heads, batch_firstTrue) self.norm2 nn.LayerNorm(dim) self.mlp nn.Sequential( nn.Linear(dim, int(dim * mlp_ratio)), nn.GELU(), nn.Linear(int(dim * mlp_ratio), dim) ) def forward(self, x): h self.norm1(x) h, _ self.attn(h, h, h) # 自注意力QKV x x h x x self.mlp(self.norm2(x)) return x class DehazeViT(nn.Module): def __init__(self, img_size256, patch_size8, embed_dim256, depth6): super().__init__() self.patch_embed PatchEmbed(img_size, patch_size, 3, embed_dim) self.blocks nn.ModuleList([ TransformerBlock(embed_dim) for _ in range(depth) ]) # 解码把 token 还原成特征图再卷积上采样 self.decoder nn.Sequential( nn.Conv2d(embed_dim, 128, 3, padding1), nn.ReLU(inplaceTrue), nn.Upsample(scale_factor2, modebilinear, align_cornersFalse), nn.Conv2d(128, 64, 3, padding1), nn.ReLU(inplaceTrue), nn.Upsample(scale_factor2, modebilinear, align_cornersFalse), nn.Conv2d(64, 32, 3, padding1), nn.ReLU(inplaceTrue), nn.Upsample(scale_factor2, modebilinear, align_cornersFalse), nn.Conv2d(32, 3, 3, padding1), nn.Sigmoid() # 输出限制在 [0,1] ) self.patch_size patch_size def forward(self, x): b, c, h, w x.shape tokens self.patch_embed(x) # B, N, C for blk in self.blocks: tokens blk(tokens) # token 还原成 2D 特征图 hp, wp h // self.patch_size, w // self.patch_size feat rearrange(tokens, b (h w) c - b c h w, hhp, wwp) out self.decoder(feat) return out逻辑说明PatchEmbed 用 stride 等于 patch_size 的卷积替代手工 unfold省显存也更快TransformerBlock 里把 LayerNorm 放在注意力之前pre-norm训练更稳这是 ViT 原论文之后的主流改法解码器用三次双线性上采样把 32x32 的特征图拉回 256x256最后 Sigmoid 保证输出在合法像素范围。参数说明patch_size 设 8 是去雾任务的常用折中——设 16 会丢细节设 4 则 token 数暴涨、显存吃不消embed_dim 256 在单张 12G 显存卡上 batch_size 能开到 8 左右depth 6 是轻量配置追求指标可以加到 12但要注意过拟合。img_size 训练时统一到 256推理时可以吃任意尺寸因为卷积和上采样都不依赖固定输入。2.3 损失函数怎么配L1、感知损失与 SSIM 的取舍只用一个 L1 损失训出来的去雾图PSNR 可能不低但肉眼看着发灰、边缘发软。我一般用组合损失L1 主回归加一项 VGG 感知损失管纹理再加一项 SSIM 损失管结构。权重上 L1 占 1.0感知损失 0.04SSIM 0.1这个比例在 RESIDE 类数据集上比较稳。import torch.nn.functional as F from pytorch_msssim import ssim def dehaze_loss(pred, gt, vgg_feat_fnNone): l1 F.l1_loss(pred, gt) ssim_loss 1 - ssim(pred, gt, data_range1.0, size_averageTrue) loss l1 0.1 * ssim_loss if vgg_feat_fn is not None: # 感知损失比较 VGG 中间层特征 pf vgg_feat_fn(pred) gf vgg_feat_fn(gt) loss loss 0.04 * F.l1_loss(pf, gf) return loss逻辑说明ssim 用 pytorch_msssim 库注意 data_range 要和你归一化方式一致输入是 [0,1] 就写 1.0写错会导致损失尺度异常。感知损失里的 vgg_feat_fn 建议取 relu3_3 层输出太浅没语义、太深丢空间细节。参数说明感知损失权重别超过 0.1否则输出会出现 VGG 特有的棋盘伪影SSIM 权重 0.1 是经验值调到 0.3 以上 PSNR 会掉。如果你的数据集配对质量一般先把 L1 训到收敛再加感知损失否则容易学偏。3. 数据准备与训练流程从 RESIDE 到自建雾图怎么喂给模型3.1 数据集选择与配对格式公开数据集里 RESIDE 是去雾绕不开的基准它分 ITS室内、OTS室外等子集配对形式是「清晰图 对应雾图」。如果你要做的是真实雾图没有 GT那就得走域适应或者用合成雾先预训练再微调。常见做法是先用 RESIDE-ITS 预训练再用少量真实配对数据 finetune。目录结构建议这样组织训练脚本按文件名配对省得写映射表dataset/ train/ clear/ 0001.png 0002.png ... hazy/ 0001.png 0002.png ... val/ clear/ ... hazy/ ...3.2 训练脚本的关键参数与启动命令from torch.utils.data import Dataset, DataLoader from PIL import Image import torchvision.transforms as T import os class DehazeDataset(Dataset): def __init__(self, root, size256): self.clear_dir os.path.join(root, clear) self.hazy_dir os.path.join(root, hazy) self.names sorted(os.listdir(self.clear_dir)) self.tf T.Compose([ T.Resize((size, size)), T.ToTensor() # 自动归一化到 [0,1] ]) def __len__(self): return len(self.names) def __getitem__(self, idx): n self.names[idx] clear Image.open(os.path.join(self.clear_dir, n)).convert(RGB) hazy Image.open(os.path.join(self.hazy_dir, n)).convert(RGB) return self.tf(hazy), self.tf(clear) # 训练循环核心 model DehazeViT().cuda() opt torch.optim.AdamW(model.parameters(), lr2e-4, weight_decay1e-4) sched torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max100) loader DataLoader(DehazeDataset(dataset/train), batch_size8, shuffleTrue, num_workers4) for epoch in range(100): model.train() for hazy, clear in loader: hazy, clear hazy.cuda(), clear.cuda() pred model(hazy) loss dehaze_loss(pred, clear) opt.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() sched.step() torch.save(model.state_dict(), fckpt/epoch_{epoch}.pth)逻辑说明ToTensor 已经把像素压到 [0,1]和网络末端 Sigmoid 输出范围对齐不用再手动除 255。梯度裁剪设 1.0 是因为 Transformer 在训练初期注意力 logits 容易爆不裁的话 loss 会突然变 NaN这个坑我踩过不止一次。参数说明lr 用 2e-4 配 AdamW比 1e-3 稳weight_decay 1e-4 抑制过拟合batch_size 8 是 12G 显存下的安全值显存够可以上 16 并同步把 lr 提到 3e-4CosineAnnealingLR 的 T_max 设成总 epoch 数别设小了否则后期 lr 归零学不动。3.3 推理与指标验证import torch from skimage.metrics import peak_signal_noise_ratio as psnr from skimage.metrics import structural_similarity as ssim_metric model.eval() with torch.no_grad(): pred model(hazy.cuda()).cpu().squeeze(0).permute(1, 2, 0).numpy() gt clear.squeeze(0).permute(1, 2, 0).numpy() print(PSNR:, psnr(gt, pred, data_range1.0)) print(SSIM:, ssim_metric(gt, pred, channel_axis2, data_range1.0))逻辑说明推理时记得 model.eval() 关掉 dropout 和 BN 统计更新否则同一张图两次结果不一样。指标计算前把 tensor 转成 HWC numpyskimage 的 SSIM 要指定 channel_axis2老版本参数名是 multichannel版本不对会报错。参数说明data_range 必须和像素范围一致[0,1] 就写 1.0写 255 指标会虚高。PSNR 到 30dB 以上、SSIM 到 0.95 以上在 RESIDE-ITS 上算及格线低于这个数先查数据配对有没有错位。4. 避坑与排查ViT 去雾训练中最容易翻车的 5 个点4.1 现象loss 前几个 epoch 就变 NaN原因Transformer 的注意力在初始化阶段输出方差大加上 lr 偏高梯度直接炸掉。解决把梯度裁剪加上clip_grad_norm_ 设 1.0lr 从 2e-4 起步warmup 前 500 步线性升到目标 lr。如果还炸检查输入有没有 NaN用 torch.isnan(hazy).any() 扫一遍。4.2 现象PSNR 涨到 28 就卡住不动原因多半是 patch_size 设太大细节恢复能力到顶了。解决把 patch_size 从 16 降到 8token 数翻四倍显存不够就减 batch_size 或加梯度累积。另一个可能是损失权重失衡感知损失占比过高先把感知损失关掉只留 L1 看能不能突破。4.3 现象输出图有规则网格状伪影原因patch 边界处特征不连续解码器上采样时把这种不连续放大了。解决在解码器末端加一个 3x3 卷积做平滑或者改用重叠 patchstride 小于 kernel_size。我一般直接在最后两层用卷积替代最近邻上采样伪影能压掉大半。4.4 现象验证集指标好但真实雾图一塌糊涂原因训练数据全是合成雾合成雾的物理模型和真实雾差得远模型学到了合成噪声的捷径。解决混入真实雾图做微调哪怕没有 GT也可以用对比学习或者伪标签。至少把合成雾的生成参数beta 散射系数、大气光 A随机化范围拉大别只用固定参数生成。4.5 现象显存明明够却报 OOM原因PyTorch 的缓存分配器碎片化或者 num_workers 开太大把内存吃满。解决设 torch.cuda.empty_cache() 在验证阶段清缓存num_workers 从 4 降到 2 试试。另外 einops 的 rearrange 在某些形状下会产生临时大张量把 patch 还原那步改成 viewpermute 能省一点。5. 进阶技巧把 ViT 去雾压到能实时跑以及怎么判断它值不值得上5.1 用窗口注意力换速度标准全局自注意力的计算量随 token 数平方增长256x256 输入、patch 8 就是 1024 个 token注意力矩阵 1024x1024推理一张图要几百毫秒。想上实时把全局注意力换成窗口注意力类似 Swin 的做法窗口设 8计算量直接降一个量级。改动很小把 TransformerBlock 里的 MultiheadAttention 换成窗口内注意力即可代价是全局建模能力弱一点但去雾任务里窗口 8 已经能覆盖大部分雾的局部一致性。def window_partition(x, window_size): B,H,W,C - B*nW, ws, ws, C B, H, W, C x.shape x x.view(B, H // window_size, window_size, W // window_size, window_size, C) windows x.permute(0, 1, 3, 2, 4, 5).contiguous() return windows.view(-1, window_size, window_size, C)逻辑说明把特征图按 window_size 切块每个窗口独立做注意力最后再 merge 回去。参数说明window_size 设 8 是速度和效果的平衡点设 4 更快但窗口间信息交换不足设 16 接近全局注意力、提速有限。5.2 判断这个方向值不值得投入的三个信号第一看你的数据里雾的分布是否全局不均匀。如果只是均匀薄雾传统方法加个轻量 CNN 就够了上 ViT 是杀鸡用牛刀。第二看你对推理延迟的容忍度。ViT 去雾在 1080Ti 上单张 256x256 大概 80 到 150 毫秒窗口版能压到 30 毫秒以内如果业务要求 10 毫秒级得考虑蒸馏到小 CNN。第三看有没有配对数据。没有配对数据的话ViT 参数量大、更容易过拟合到合成雾的分布反而不如小模型稳。5.3 一个我常用的验证习惯每次改完结构我不急着跑完整训练先拿 50 张图过拟合一遍。如果模型连这 50 张都压不到 PSNR 35 以上说明结构或损失有问题不用浪费卡时跑全量。这个习惯帮我省下的 GPU 时间比任何调参技巧都实在。另外权重文件别只存最后一个 epoch每 10 个 epoch 存一次去雾任务后期指标会波动最优模型往往不在最后一轮。希望帮到你。本文还有配套的精品资源点击获取
返回列表