多彩编程 多彩编程MZPH · CODE BLOG
ARTICLE DETAIL

文章详情

深耕前端与后端开发技术的一线实战笔记与踩坑复盘。

PyTorch对偶GAN图像去雾实战:从环境搭建到训练调优

PyTorch对偶GAN图像去雾实战:从环境搭建到训练调优 简介这份资源是面向计算机相关专业毕业设计、课程设计及期末作业场景的PyTorch实战项目核心任务是用对偶生成对抗网络完成图像去雾。项目由生成器与判别器双网络协同训练配套训练、预测、参数解析、数据加载与可视化等模块适合已具备Python与深度学习基础、希望用完整项目提升工程能力的学习者。压缩包共31个文件约21.31MB包含10个py源码、6个png与5个jpg效果图、4个zbak备份、2个pkl模型权重以及license、md说明等源码与文档均带逐行注释便于理解网络结构与训练流程。目前已有43人学习。读者可据此掌握对偶GAN去雾的完整实现思路参考模型权重与预测脚本快速复现效果并借助备份文件与说明文档排查训练中的常见问题也可作为毕业设计或课程作业的参考方案。1. 从一张雾天照片说起这套 PyTorch 对偶 GAN 去雾源码到底能跑出什么雾天拍出来的照片最直观的感受就是对比度塌了、远处细节糊成一团做目标检测或者视频监控时模型基本等于瞎了一半。传统暗通道先验那套方法我早些年也用过遇到天空区域或者大面积白墙就翻车颜色偏得离谱。这套基于 PyTorch 的对偶生成对抗网络图像去雾系统走的是另一条路用两个生成器分别学「有雾到无雾」和「无雾到有雾」两个方向的映射再配两个判别器做对抗约束形成闭环。它解决的核心问题是——在没有成对训练数据的情况下也能训出一个能用的去雾模型。适合谁正在做 PyTorch 生成对抗网络方向毕业设计的学生或者手头有雾天数据集、想快速搭一套可训练可推理 pipeline 的工程师。源码包里把数据加载、模型定义、训练循环、推理脚本都拆开了不是那种一个文件糊到底的写法改起来心里有底。2. 对偶 GAN 去雾的原理拆解与 PyTorch 环境搭建2.1 为什么选对偶结构而不是单生成器单生成器去雾的思路很直接输入有雾图输出无雾图判别器判断输出像不像真实无雾图。但这里有个致命问题——没有成对的有雾无雾数据时生成器很容易把所有输入都映射成同一张「安全」的输出比如整体偏灰、细节全丢判别器也分不出来。对偶 GAN 的核心在于 Cycle Consistency Loss有雾图经过生成器 G 变成无雾图再经过生成器 F 变回有雾图这个来回的结果要和原图尽可能一致。反过来也一样。这样一来生成器不能随便糊弄必须保留足够的内容信息才能循环回来。具体到去雾任务两个生成器的分工是G_AB有雾域 A → 无雾域 BG_BA无雾域 B → 有雾域 AD_A判别输入是真实有雾图还是 G_BA 生成的假有雾图D_B判别输入是真实无雾图还是 G_AB 生成的假无雾图损失函数由三部分组成对抗损失、循环一致性损失、身份损失。身份损失是我个人觉得比较关键的一项它约束 G_AB(B) ≈ B也就是无雾图过一遍生成器不应该被改变太多防止生成器过度发挥。2.2 PyTorch 环境搭建版本对应与 CUDA 选择环境这块踩过的坑太多了先给一个我验证过的组合组件版本说明Python3.8 ~ 3.103.11 以上有些旧版 torchvision 轮子不全PyTorch1.12 ~ 2.02.x 也能跑但要注意torch.cuda.amp的 API 变化CUDA11.3 / 11.6 / 11.8跟 PyTorch 官方对应表走torchvision与 PyTorch 匹配别单独 pip install容易版本打架其他numpy, pillow, tqdm, tensorboard常规依赖安装命令按 CUDA 版本选比如 CUDA 11.8# CUDA 11.8 对应的 PyTorch 安装 pip install torch2.0.1 torchvision0.15.2 --index-url https://download.pytorch.org/whl/cu118 # 如果只是 CPU 跑推理用这个 pip install torch2.0.1 torchvision0.15.2 --index-url https://download.pytorch.org/whl/cpu逻辑说明--index-url指定 PyTorch 官方轮子源避免从默认源拉到 CPU 版本。参数上cu118对应 CUDA 11.8如果你装的是 CUDA 11.6 就换成cu116。装完用下面这段验证import torch print(torch.__version__) # 确认 PyTorch 版本 print(torch.cuda.is_available()) # True 才说明 GPU 可用 print(torch.version.cuda) # 确认 CUDA 版本如果is_available()返回 False先别急着重装八成是驱动版本和 CUDA 运行时不匹配。用nvidia-smi看驱动支持的 CUDA 上限再决定装哪个版本的 PyTorch。提示Anaconda 环境下建议单独建虚拟环境conda create -n dehaze python3.9别在 base 里折腾后面依赖冲突会让人崩溃。2.3 源码目录结构与模块职责拿到源码包后先别急着跑train.py。花五分钟把目录过一遍后面改代码能省很多时间。典型结构是这样的dehaze-gan/ ├── data/ # 数据集存放 │ ├── trainA/ # 有雾图 │ ├── trainB/ # 无雾图 │ └── testA/ # 测试用有雾图 ├── models/ │ ├── generator.py # 生成器网络定义 │ ├── discriminator.py # 判别器网络定义 │ └── cycle_gan.py # 对偶 GAN 整体封装 ├── utils/ │ ├── dataset.py # 数据加载与预处理 │ ├── visualizer.py # 训练过程可视化 │ └── metrics.py # PSNR / SSIM 计算 ├── train.py # 训练入口 ├── test.py # 推理入口 └── config.yaml # 超参数配置生成器一般用 ResNet 风格的编码器-解码器结构中间堆 6 到 9 个残差块。判别器用 PatchGAN输出一个 N×N 的 patch 判别结果而不是单个标量。这个设计对去雾任务比较友好因为它关注局部纹理是否真实而不是整图统计量。3. 训练流程实操数据准备、超参设置与损失曲线观察3.1 数据集组织与预处理对偶 GAN 不需要成对数据但需要两个域的图片分别放在trainA和trainB里。有雾图可以来自 RESIDE 的合成雾图也可以自己用大气散射模型合成。无雾图用清晰的自然图像即可。数量上不要求一一对应但两个域的数量级别差太多否则判别器会偏向数量多的那一域。预处理这块源码里一般会做随机裁剪到 256×256、随机水平翻转、归一化到 [-1, 1]。注意归一化范围要和生成器最后一层的激活函数匹配用tanh就归到 [-1, 1]用sigmoid就归到 [0, 1]这个对错了训练直接不收敛。# utils/dataset.py 中的关键预处理逻辑 from torchvision import transforms transform transforms.Compose([ transforms.Resize((286, 286)), # 先放大再随机裁剪增加多样性 transforms.RandomCrop(256), # 裁剪到 256x256 transforms.RandomHorizontalFlip(), # 水平翻转增强 transforms.ToTensor(), # 转成 tensor范围 [0, 1] transforms.Normalize((0.5, 0.5, 0.5), # 归一化到 [-1, 1] (0.5, 0.5, 0.5)) ])逻辑说明Resize到 286 再RandomCrop到 256是 CycleGAN 原论文的做法目的是让每次裁剪的位置有变化。Normalize的均值和标准差都设 0.5把 [0, 1] 映射到 [-1, 1]。如果你的生成器输出层用的是Sigmoid这里就要改成Normalize((0, 0, 0), (1, 1, 1))保持范围一致。3.2 超参数配置与训练命令配置文件里几个关键参数直接决定训练能不能收敛# config.yaml 关键字段 lr_G: 0.0002 # 生成器学习率 lr_D: 0.0002 # 判别器学习率 beta1: 0.5 # Adam 优化器的 beta1 beta2: 0.999 # Adam 优化器的 beta2 lambda_cycle: 10.0 # 循环一致性损失权重 lambda_identity: 5.0 # 身份损失权重 batch_size: 1 # 对偶 GAN 显存占用大通常设 1 n_epochs: 200 # 总训练轮数 decay_epoch: 100 # 学习率开始线性衰减的轮数lambda_cycle设 10 是 CycleGAN 原论文的默认值去雾任务上我试过 5 到 15 之间10 比较稳。lambda_identity设 5 能有效抑制生成器乱改颜色。batch_size设 1 是因为对偶 GAN 要同时跑两个生成器和两个判别器显存吃紧Instance Normalization 在 batch size 为 1 时也能工作。启动训练python train.py --config config.yaml --dataroot ./data --name dehaze_exp1训练过程中重点看两个东西一是 TensorBoard 里的损失曲线G_A 和 G_B 的损失应该在震荡中缓慢下降如果某一方持续为 0 或者爆炸说明对抗失衡了二是每轮保存的生成样本看有雾图过 G_AB 之后是不是真的变清晰了而不是变成一片灰。3.3 损失函数代码解析对偶 GAN 的损失计算是整个训练的核心源码里一般封装在cycle_gan.py的compute_losses方法里# models/cycle_gan.py 损失计算核心逻辑 import torch.nn.functional as F def compute_losses(self, real_A, real_B): # 生成器前向 fake_B self.G_AB(real_A) # 有雾 - 无雾 fake_A self.G_BA(real_B) # 无雾 - 有雾 rec_A self.G_BA(fake_B) # 循环回来 rec_B self.G_AB(fake_A) idt_A self.G_BA(real_A) # 身份映射 idt_B self.G_AB(real_B) # 对抗损失最小二乘 GAN 损失比 BCE 更稳 loss_GAN_AB F.mse_loss(self.D_B(fake_B), torch.ones_like(self.D_B(fake_B))) loss_GAN_BA F.mse_loss(self.D_A(fake_A), torch.ones_like(self.D_A(fake_A))) # 循环一致性损失 loss_cycle_A F.l1_loss(rec_A, real_A) * self.lambda_cycle loss_cycle_B F.l1_loss(rec_B, real_B) * self.lambda_cycle # 身份损失 loss_idt_A F.l1_loss(idt_A, real_A) * self.lambda_cycle * self.lambda_identity loss_idt_B F.l1_loss(idt_B, real_B) * self.lambda_cycle * self.lambda_identity # 生成器总损失 loss_G loss_GAN_AB loss_GAN_BA loss_cycle_A loss_cycle_B loss_idt_A loss_idt_B return loss_G逻辑说明对抗损失用mse_loss而不是binary_cross_entropy这是 LSGAN 的做法梯度更平滑训练早期不容易崩。循环损失和身份损失都用 L1因为 L1 对图像重建的约束比 L2 更锐利不容易糊。参数上lambda_cycle和lambda_identity的乘积关系要注意源码里如果写成lambda_cycle * lambda_identity那实际身份损失权重就是 50偏大我一般会把lambda_identity单独设成 0.5 到 1.0 之间。4. 推理与效果验证从单张图到批量测试的完整链路4.1 单张图像推理脚本训练完之后最关心的就是拿一张新雾图进去出来的结果到底能不能看。推理脚本一般长这样# test.py 单张推理核心逻辑 import torch from PIL import Image from torchvision import transforms from models.generator import ResnetGenerator def dehaze_single(image_path, weight_path, output_path): device torch.device(cuda if torch.cuda.is_available() else cpu) # 加载生成器结构并载入权重 G ResnetGenerator(input_nc3, output_nc3, n_blocks9) G.load_state_dict(torch.load(weight_path, map_locationdevice)) G.to(device).eval() # 预处理保持和训练时一致 transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) img Image.open(image_path).convert(RGB) input_tensor transform(img).unsqueeze(0).to(device) # 推理并反归一化 with torch.no_grad(): output G(input_tensor) output output.squeeze(0).cpu() output (output * 0.5 0.5).clamp(0, 1) # 反归一化到 [0, 1] out_img transforms.ToPILImage()(output) out_img.save(output_path) dehaze_single(testA/foggy_01.jpg, checkpoints/G_AB_latest.pth, result_01.png)逻辑说明eval()必须调用否则 BatchNorm 或 Dropout 层会按训练模式跑结果不稳定。torch.no_grad()关掉梯度计算省显存也提速。反归一化那步output * 0.5 0.5对应训练时的Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))如果训练时归一化参数不同这里要同步改。n_blocks9要和训练时的生成器配置一致否则load_state_dict会报 key 不匹配。4.2 批量测试与指标计算单张看着不错不代表整体好批量跑一遍测试集算 PSNR 和 SSIM 才有说服力。源码里utils/metrics.py一般会封装好# utils/metrics.py 指标计算 import torch import torch.nn.functional as F from skimage.metrics import structural_similarity as ssim import numpy as np def compute_psnr(pred, target): # pred 和 target 都是 [0, 1] 范围的 tensor mse F.mse_loss(pred, target) if mse 0: return float(inf) return 10 * torch.log10(1.0 / mse).item() def compute_ssim(pred, target): # 转成 numpy 算 SSIM pred_np pred.squeeze(0).permute(1, 2, 0).cpu().numpy() target_np target.squeeze(0).permute(1, 2, 0).cpu().numpy() return ssim(pred_np, target_np, channel_axis2, data_range1.0)逻辑说明PSNR 基于 MSE 计算对亮度误差敏感SSIM 关注结构相似性更贴近人眼感受。两个指标要一起看PSNR 高但 SSIM 低说明可能整体亮度对了但结构糊了。参数上data_range1.0对应输入范围 [0, 1]如果输入是 [0, 255] 就要改成 255。批量测试时把每张图的指标存到 CSV 里最后算平均值和标准差标准差大说明模型在某些场景下不稳定。注意合成雾图和真实雾图的指标没有可比性。如果你的测试集是合成的PSNR 可能到 25 以上换成真实雾图PSNR 掉到 15 以下也正常别拿合成集的指标去吹真实场景的效果。4.3 效果对比与失败案例分析跑完批量测试挑几张典型图出来对比。我一般会分三类看第一类薄雾场景远处有轻微灰蒙感。这种对偶 GAN 处理起来比较稳去雾后对比度提升明显颜色也不会偏太多。第二类浓雾场景能见度很低。这种容易翻车生成器可能把雾去掉的同时也把远处细节一起抹了输出看起来像过度锐化或者油画效果。原因是循环一致性损失在浓雾下约束不够强生成器倾向于输出「安全」的平滑结果。第三类天空区域占比大的图。天空在雾天里亮度高、纹理少判别器很难判断去雾后的天空是否真实生成器容易把天空处理成奇怪的色块。常见做法是在损失里加一个天空区域的权重掩码或者训练时多喂一些带天空的样本。失败案例别删留着分析。我习惯把失败图按场景分类看是模型容量不够还是数据分布没覆盖到。如果是某一类场景反复失败优先补数据而不是调网络结构。5. 避坑与排查训练不收敛、显存爆炸、颜色偏移的实战记录5.1 生成器损失持续为 0 或判别器完全碾压现象训练日志里loss_D很快降到接近 0loss_G居高不下生成的图要么全黑要么全白。原因判别器太强生成器梯度消失。对偶 GAN 里两个判别器同时训练如果判别器学习率偏高或者更新次数过多就会压倒生成器。解决把lr_D降到lr_G的一半或者每训练 2 次生成器才更新 1 次判别器。另外检查对抗损失是不是用了 BCE换成 LSGAN 的 MSE 损失会稳很多。如果已经崩了别接着训回滚到最近一个生成效果还行的 checkpoint 重新调参。5.2 显存爆炸batch size 设 1 还是 OOM现象RuntimeError: CUDA out of memory即使batch_size1也报错。原因对偶 GAN 同时加载两个生成器、两个判别器加上中间特征图和梯度显存占用是单生成器的 4 倍左右。如果输入分辨率是 512×512 而不是 256×256显存直接翻四倍。解决先把输入裁剪到 256×256 跑通再考虑加大。开启混合精度训练能省 30% 到 40% 显存# 混合精度训练核心代码 from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for real_A, real_B in dataloader: optimizer_G.zero_grad() with autocast(): # 前向用半精度 loss_G model.compute_losses(real_A, real_B) scaler.scale(loss_G).backward() # 梯度缩放防止下溢 scaler.step(optimizer_G) scaler.update()逻辑说明autocast让前向计算自动用 float16GradScaler对损失做缩放再反向传播避免半精度下梯度太小变成 0。参数上不需要手动调GradScaler会自动管理缩放因子。注意optimizer.step()要用scaler.step()替代否则缩放没生效。5.3 去雾后颜色整体偏移现象输出图确实清晰了但整体偏蓝或者偏黄和真实无雾图色温对不上。原因身份损失权重太低生成器在去雾的同时把颜色也改了。或者训练集里无雾图的色温分布和有雾图差异太大模型学到了错误的颜色映射。解决把lambda_identity从 5 提到 10增强身份约束。同时检查trainB里的无雾图是不是色温统一如果一半冷色调一半暖色调模型会懵。常见做法是先对训练集做白平衡统一再喂进去训。5.4 推理结果和训练时看到的样本差距大现象训练过程中保存的生成样本看着不错但用test.py推理同一张图结果差很多。原因推理时的预处理和训练时不一致。比如训练用了RandomCrop到 256推理时直接整图输入分辨率不匹配导致生成器行为异常。或者归一化参数写错了训练用 [-1,1]推理忘了反归一化。解决把推理的预处理单独写成一个函数和训练时的transform保持完全一致只去掉随机增强部分。输入分辨率要么和训练一致要么用Resize到训练分辨率再推理。反归一化那步一定要检查输出图发灰或者过曝基本都是这里的问题。5.5 训练到一半 loss 突然 NaN现象前 50 轮都正常突然某一轮 loss 变成 NaN之后所有输出都是噪声。原因学习率偏高导致梯度爆炸或者某张异常图片全黑、全白、损坏进入训练。对偶 GAN 的循环损失里如果有除法操作遇到零值也会产生 NaN。解决加梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)把梯度范数限制在 1 以内。同时在数据加载时过滤掉尺寸过小或者像素值全零的图片。如果已经 NaN 了从最近的正常 checkpoint 恢复把学习率乘 0.5 再继续。6. 进阶技巧用感知损失和注意力机制把去雾效果再拉一档基础版对偶 GAN 跑通之后如果效果还差口气可以往上叠两个东西感知损失和注意力机制。这两个我在实际项目里都试过提升是实打实的但要注意叠加顺序和权重。先说感知损失。L1 和 L2 损失关注的是像素级差异但人眼对结构、纹理、边缘更敏感。感知损失的做法是把生成图和真实图同时送进一个预训练的 VGG 网络取中间某几层的特征图算 L1 距离。这样约束的是「语义层面像不像」而不是「每个像素差多少」。代码大概长这样# 感知损失实现 import torchvision.models as models class PerceptualLoss(torch.nn.Module): def __init__(self): super().__init__() vgg models.vgg16(pretrainedTrue).features # 取前 16 层覆盖浅层纹理和深层语义 self.slice torch.nn.Sequential(*list(vgg[:16])).eval() for p in self.parameters(): p.requires_grad False # VGG 不参与训练 def forward(self, pred, target): # VGG 输入需要归一化到 ImageNet 统计量 mean torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(pred.device) std torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(pred.device) pred_norm (pred - mean) / std target_norm (target - mean) / std feat_pred self.slice(pred_norm) feat_target self.slice(target_norm) return F.l1_loss(feat_pred, feat_target)逻辑说明vgg[:16]取的是 VGG16 前 16 层包含 3 个卷积块能同时捕捉浅层边缘和深层语义。requires_gradFalse冻结 VGG 参数只把它当特征提取器用。输入归一化必须用 ImageNet 的均值和标准差因为 VGG 是在 ImageNet 上预训练的不归一化特征会偏。感知损失的权重一般设 0.1 到 0.5 之间太大反而会让输出过度平滑。再说注意力机制。去雾任务里雾的分布是不均匀的远处雾浓、近处雾薄如果生成器对所有区域一视同仁浓雾区域去不干净薄雾区域又过度处理。加一个通道注意力模块让网络自己学哪些通道对去雾更重要# 通道注意力模块 class ChannelAttention(torch.nn.Module): def __init__(self, channels, reduction16): super().__init__() self.avg_pool torch.nn.AdaptiveAvgPool2d(1) self.fc torch.nn.Sequential( torch.nn.Linear(channels, channels // reduction), torch.nn.ReLU(inplaceTrue), torch.nn.Linear(channels // reduction, channels), torch.nn.Sigmoid() ) def forward(self, x): b, c, _, _ x.size() y self.avg_pool(x).view(b, c) y self.fc(y).view(b, c, 1, 1) return x * y # 通道加权逻辑说明AdaptiveAvgPool2d(1)把每个通道的空间信息压成一个标量再通过两个全连接层学出每个通道的权重最后乘回原特征图。reduction16是压缩比通道数少的时候可以设 8。这个模块插在生成器的残差块之间不改变特征图尺寸即插即用。两个技巧叠加的时候我一般先加感知损失跑 50 轮看效果稳定了再加注意力模块。同时加容易让训练初期不稳定损失震荡得厉害。权重上感知损失 0.2、注意力模块不加额外权重让网络自己学。验证方法还是老一套批量跑测试集算 PSNR 和 SSIM和基础版对比两个指标都涨了才说明有效只涨一个可能是过拟合了。从那以后我每次改损失函数或者加模块都强制走一遍「基础版指标 → 改一版 → 对比指标」的流程不靠肉眼感觉。希望帮到你。本文还有配套的精品资源点击获取
返回列表