简介这份资源是面向深度学习研究者与图像处理方向学习者的PyTorch版DehazeNet图像去雾实现提供从网络结构定义、训练流程到推理演示的完整代码链路并附带已训练好的预训练权重可直接加载使用也便于在现有架构上做扩展实验。压缩包共21个文件约114KB以9个Python脚本和4个MATLAB文件为主前者覆盖训练、验证与预处理等环节后者用于引导滤波等辅助处理另有2个pth权重文件、若干备份文件及说明文档整体结构紧凑、便于快速上手。目前已有86人学习关注。借助其中的室内与室外两套预训练模型、训练与验证脚本以及演示程序读者能够直接复现去雾效果并对照代码理解数据准备、模型训练与推理的完整流程适合具备一定深度学习基础、希望深入掌握去雾算法实现细节的研究人员参考。1. DehazeNet 去雾到底值不值得做从一张灰蒙蒙的照片说起拍过户外照片的人都有体会明明肉眼看着通透的天气出片却像蒙了一层灰纱远处建筑轮廓发白树叶颜色发闷。这不是相机坏了是大气散射在作怪——悬浮颗粒把光线散射掉成像传感器收到的对比度和色彩都被稀释了。传统做法靠暗通道先验、大气光估计那一套参数调起来玄学得很换一批图就得重来。DehazeNet 的思路不一样它用卷积网络直接学「有雾图到无雾图」的映射关系把去雾当成一个端到端的回归问题。这篇要讲的就是基于 PyTorch 把 DehazeNet 从零搭起来、训练、加载预训练模型推理的完整路径顺带把环境搭建、参数设置、翻车点都摊开说。适合已经会一点 PyTorch、想拿去雾练手或落地到实际图像增强流程里的同学新手照着步骤也能跑通熟手可以重点看后面的边界和坑。2. DehazeNet 的网络结构与 PyTorch 实现细节2.1 为什么 DehazeNet 不是简单堆卷积DehazeNet 的核心设计动机来自大气散射模型I(x) J(x)t(x) A(1 - t(x))其中 I 是有雾图J 是无雾图t 是透射率A 是大气光。传统方法要分别估计 t 和 ADehazeNet 则让网络直接预测透射率图再用大气光做还原。它的结构有几个关键点第一层用多尺度卷积核比如 5x5、7x7 并行提取不同感受野的特征这比单一尺寸卷积更能捕捉雾的局部浓度变化中间用 Maxout 激活而不是 ReLU因为 Maxout 能拟合更复杂的非线性对透射率这种连续值回归更友好最后用 BReLU双边 ReLU把输出限制在 [0,1] 区间因为透射率物理上就在这个范围。这些设计不是拍脑袋是作者针对去雾任务特性做的取舍。用 PyTorch 实现时Maxout 没有现成层得自己写BReLU 也要自定义。下面先把网络骨架搭出来。import torch import torch.nn as nn import torch.nn.functional as F class Maxout(nn.Module): Maxout 激活把通道分组每组取最大值 def __init__(self, in_channels, out_channels, groups2): super().__init__() self.groups groups self.linear nn.Conv2d(in_channels, out_channels * groups, 1) def forward(self, x): out self.linear(x) # 按 groups 维度分组取 max out out.view(out.size(0), self.groups, -1, out.size(2), out.size(3)) return out.max(dim1)[0] class BReLU(nn.Module): 双边 ReLU下界 0上界 1 def forward(self, x): return torch.clamp(x, 0.0, 1.0) class DehazeNet(nn.Module): def __init__(self): super().__init__() # 多尺度特征提取并行不同卷积核 self.conv1_1 nn.Conv2d(3, 16, 5, padding2) self.conv1_2 nn.Conv2d(3, 16, 7, padding3) self.conv1_3 nn.Conv2d(3, 16, 3, padding1) # Maxout 融合 self.maxout1 Maxout(48, 16, groups2) # 中间层 self.conv2 nn.Conv2d(16, 16, 3, padding1) self.maxout2 Maxout(16, 16, groups2) self.conv3 nn.Conv2d(16, 16, 3, padding1) self.maxout3 Maxout(16, 16, groups2) # 输出透射率图单通道 self.conv4 nn.Conv2d(16, 1, 3, padding1) self.brelu BReLU() def forward(self, x): x1 self.conv1_1(x) x2 self.conv1_2(x) x3 self.conv1_3(x) x torch.cat([x1, x2, x3], dim1) x self.maxout1(x) x self.maxout2(self.conv2(x)) x self.maxout3(self.conv3(x)) t self.brelu(self.conv4(x)) return t这段代码里Maxout 的实现用 1x1 卷积把通道数扩大 groups 倍再 reshape 后取 max等价于原论文的 Maxout 操作。参数 groups2 是常用值调大能增加非线性但显存和计算量上升。BReLU 用 clamp 实现简单直接。多尺度卷积并行后 concat再进 Maxout 融合这是 DehazeNet 区别于普通 CNN 的关键。注意 conv1 的三个分支输出通道都是 16concat 后是 48Maxout 再压回 16。如果显存吃紧可以把通道数减半但透射率估计精度会掉一点这个后面避坑章会细说。2.2 用 PyTorch 搭训练流程数据、损失、优化器怎么配网络有了接下来是训练。去雾数据集常见的有 RESIDE、NYU-Depth 的合成雾图、或者自己用大气散射模型合成。合成时透射率 t 和大气光 A 是已知的所以可以直接监督透射率图损失用 MSE 就行。但实际中更稳的做法是监督最终去雾结果 J因为透射率的小误差经过还原公式会放大。我一般会两个都加透射率损失权重 0.5去雾结果损失权重 1.0。优化器用 Adam学习率 1e-3每 20 个 epoch 衰减 0.5。batch size 根据显存来8G 卡用 8 或 16 都行。下面是一个训练循环的骨架。from torch.utils.data import DataLoader, Dataset import torch.optim as optim class HazeDataset(Dataset): 假设数据目录下 hazy/ 和 clear/ 成对存放 def __init__(self, hazy_dir, clear_dir, transformNone): self.hazy_files sorted(os.listdir(hazy_dir)) self.hazy_dir hazy_dir self.clear_dir clear_dir self.transform transform def __len__(self): return len(self.hazy_files) def __getitem__(self, idx): name self.hazy_files[idx] hazy Image.open(os.path.join(self.hazy_dir, name)).convert(RGB) clear Image.open(os.path.join(self.clear_dir, name)).convert(RGB) if self.transform: hazy self.transform(hazy) clear self.transform(clear) return hazy, clear # 数据加载 transform transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor(), ]) dataset HazeDataset(data/hazy, data/clear, transform) loader DataLoader(dataset, batch_size8, shuffleTrue, num_workers4) # 模型、优化器、损失 device torch.device(cuda if torch.cuda.is_available() else cpu) model DehazeNet().to(device) optimizer optim.Adam(model.parameters(), lr1e-3) scheduler optim.lr_scheduler.StepLR(optimizer, step_size20, gamma0.5) mse nn.MSELoss() # 训练循环 for epoch in range(100): model.train() total_loss 0 for hazy, clear in loader: hazy, clear hazy.to(device), clear.to(device) t_pred model(hazy) # 用真实透射率监督合成数据里 t 已知这里简化用 1 - 暗通道近似 # 实际训练建议直接用 clear 图监督去雾结果 J_pred (hazy - 0.8 * (1 - t_pred)) / (t_pred 1e-6) # A 取 0.8 近似 J_pred torch.clamp(J_pred, 0, 1) loss mse(J_pred, clear) 0.5 * mse(t_pred, torch.ones_like(t_pred) * 0.5) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() scheduler.step() print(fEpoch {epoch}, Loss: {total_loss / len(loader):.4f})这里有几个参数要说明大气光 A 在合成数据里通常是已知的但推理时不知道常见做法是用暗通道最亮像素估计或者直接取 0.8~0.95 的经验值。学习率 1e-3 是 Adam 的常用起点如果 loss 震荡就降到 5e-4。num_workers 在 Windows 下有时会报错设成 0 就行。损失函数里透射率监督那项如果数据集没有真实 t可以去掉只监督 J。训练 100 epoch 在 RESIDE 子集上大概能收敛但具体要看数据量。显存不够就减 batch size 或把图像 resize 到 128。3. 预训练模型怎么用加载、推理与效果验证3.1 加载预训练权重并跑通单张图推理拿到预训练模型后第一件事是确认权重文件和网络结构对得上。常见坑是 state_dict 的 key 不匹配比如保存时用了 DataParallel 导致 key 前面多出 module. 前缀。加载时用 strictFalse 能跳过不匹配的层但最好还是手动去掉前缀。下面是一个完整的推理脚本。import torch from PIL import Image import torchvision.transforms as transforms def load_pretrained(model_path, device): model DehazeNet().to(device) state_dict torch.load(model_path, map_locationdevice) # 处理 DataParallel 保存的权重 new_state_dict {} for k, v in state_dict.items(): if k.startswith(module.): new_state_dict[k[7:]] v else: new_state_dict[k] v model.load_state_dict(new_state_dict, strictFalse) model.eval() return model def dehaze_image(model, img_path, device, A0.85): img Image.open(img_path).convert(RGB) transform transforms.Compose([ transforms.ToTensor(), ]) input_tensor transform(img).unsqueeze(0).to(device) with torch.no_grad(): t model(input_tensor) # 还原公式 hazy input_tensor J (hazy - A * (1 - t)) / (t 1e-6) J torch.clamp(J, 0, 1) # 转回 PIL out transforms.ToPILImage()(J.squeeze(0).cpu()) return out device torch.device(cuda if torch.cuda.is_available() else cpu) model load_pretrained(dehazenet_pretrained.pth, device) result dehaze_image(model, test_hazy.jpg, device) result.save(test_clear.jpg)这段代码里A 取 0.85 是经验值实际可以根据图像最亮区域自适应估计。t 1e-6 是防止除零。推理时不需要梯度用 torch.no_grad() 省显存。如果结果偏暗或偏亮调 A 就行A 越大去雾越强但容易过曝。注意输入图像尺寸最好和训练时一致不一致的话透射率图会有尺度偏差常见做法是 resize 到 256x256 再还原回去。3.2 用 PSNR 和 SSIM 验证去雾效果光看肉眼看不出量化差异得用指标。PSNR 和 SSIM 是去雾领域最常用的两个。PSNR 衡量像素误差SSIM 衡量结构相似性。计算时要注意图像范围是 [0,1] 还是 [0,255]别搞混。下面是一个计算脚本。import numpy as np from skimage.metrics import peak_signal_noise_ratio, structural_similarity def evaluate(clear_img, dehazed_img): # 转 numpy范围 [0,1] clear np.array(clear_img).astype(np.float32) / 255.0 dehazed np.array(dehazed_img).astype(np.float32) / 255.0 psnr peak_signal_noise_ratio(clear, dehazed, data_range1.0) ssim structural_similarity(clear, dehazed, multichannelTrue, data_range1.0) return psnr, ssim # 假设 clear_img 和 dehazed_img 都是 PIL Image psnr, ssim evaluate(clear_img, dehazed_img) print(fPSNR: {psnr:.2f}, SSIM: {ssim:.4f})参数说明data_range1.0 对应 [0,1] 范围如果图像是 [0,255] 就改成 255。multichannelTrue 处理 RGB 三通道。一般去雾任务 PSNR 能到 20 以上、SSIM 到 0.85 以上就算不错但具体看数据集难度。合成雾图指标会高些真实雾图因为没有 ground truth只能靠主观评价或无参考指标。4. 避坑与排查DehazeNet 训练推理中的 5 个血泪教训4.1 现象训练 loss 不降反升输出全黑或全白原因BReLU 把输出限制在 [0,1]但如果初始化权重太大透射率一开始就饱和到 0 或 1梯度传不回去。另外学习率设太高也会导致震荡。解决权重初始化用 kaiming_normal_偏置设 0。学习率从 1e-4 开始试别一上来就 1e-3。如果已经饱和加个梯度裁剪 torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)。4.2 现象推理结果颜色失真偏蓝或偏黄原因大气光 A 估计不准。A 取固定值 0.85 对某些图合适对另一些就不行。偏蓝通常是 A 估低了偏黄是 A 估高了。解决用暗通道先验估计 A取暗通道最亮 0.1% 像素对应原图最亮区域均值作为 A。代码里加个自适应估计别死用固定值。4.3 现象GPU 显存溢出batch size 降到 1 还报错原因多尺度卷积并行那层concat 后通道数 48中间特征图如果分辨率是 512x512显存占用是 256x256 的 4 倍。解决训练时 resize 到 256x256 或更小推理时再恢复原尺寸。或者把第一层多尺度分支减到两个通道数从 16 降到 8。显存 8G 以下建议 batch size 不超过 4。4.4 现象加载预训练模型报 key 不匹配missing keys 一大堆原因保存权重时用了 nn.DataParallelkey 前面多了 module. 前缀或者网络结构改过层名对不上。解决加载时手动去前缀用 strictFalse 先跑通再打印 missing_keys 和 unexpected_keys 看差在哪。如果是结构改了要么改回去要么只加载匹配的层。4.5 现象去雾后图像出现网格状伪影或块效应原因卷积 padding 方式不对或者用了 stride 导致棋盘效应。DehazeNet 里如果某层 stride 设成 2上采样时容易出网格。解决所有卷积层 padding 设成 kernel_size//2保持分辨率不变。别用 stride 下采样DehazeNet 原设计就是全分辨率。如果非要下采样用双线性插值上采样别用转置卷积。5. 进阶技巧把 DehazeNet 转成 ONNX 并做推理加速训练完的 PyTorch 模型部署时往往要转 ONNX方便在 C 或 TensorRT 里跑。转的时候注意动态轴设置batch 和宽高都设成动态这样一张图和多张图都能用。下面是一个转换脚本。import torch.onnx model DehazeNet().to(cpu) model.load_state_dict(torch.load(dehazenet_pretrained.pth, map_locationcpu)) model.eval() dummy_input torch.randn(1, 3, 256, 256) torch.onnx.export( model, dummy_input, dehazenet.onnx, input_names[input], output_names[transmission], dynamic_axes{ input: {0: batch, 2: height, 3: width}, transmission: {0: batch, 2: height, 3: width} }, opset_version11 )opset_version 用 11 比较稳再高有些推理引擎不支持。dynamic_axes 里把 batch 和空间维度都设成动态这样 512x512 的图也能直接推理不用改模型。转完后用 onnxruntime 验证一下输出和 PyTorch 是否一致误差在 1e-4 以内就算成功。如果要做 TensorRT 加速注意 Maxout 里的 view 和 max 操作可能被融合实际测一下速度提升一般能快 1.5 到 2 倍。我自己的习惯是每次改完网络结构先跑一遍 ONNX 导出确认没有不支持的算子再继续训练。这个后悔药提前吃比部署时才发现问题强。希望帮到你。本文还有配套的精品资源点击获取