★ FEATURED ARTICLE 简介CycleGAN的PyTorch实现面向计算机视觉与生成对抗网络方向的学习者用于完成无配对图像到图像的转换任务例如苹果与橙子互转、风格迁移等。资源包共18个文件以7个Python脚本为核心涵盖模型定义、训练器、数据工具函数等完整代码另含4个pt预训练权重可直接进行推理或继续训练2张损失曲线图帮助观察生成器与判别器的收敛过程2张转换效果示例图可用于效果对照还有1个ipynb演示笔记和项目说明文档便于快速上手压缩包总大小仅7.69MB目录组织清晰适合在本地复现和二次开发。目前已有3004人学习下载。借助预训练权重与代码读者可快速搭建环境体验从数据加载到图像转换的完整流程同时可基于现有实现调整循环一致性损失、替换生成器骨干或数据集进一步探索域适应、图像编辑等应用从模型初始化、损失计算到权重保存的完整链路均有清晰实现适合具备PyTorch基础、希望深入理解CycleGAN内部机制的中高级开发者。1. CycleGAN是什么先搞懂这个项目解决了什么问题如果你手上有一批没有配对样本的图像比如一千张夏天的街景和一千张冬天的街景它们之间不存在任何一张“同机位同角度”的对应关系但你又想把夏天的照片整体转换成冬天的质感CycleGAN 就是目前最顺手的解法。这个项目用 PyTorch 从零实现了 CycleGAN 的完整训练链路两个生成器、两个判别器、循环一致性损失、身份损失以及一套可以直接跑起来的命令行工具。你只需要把数据整理成两个文件夹就能训练出一个把 A 域转换成 B 域的模型反过来也能生成。它适合两类人一类是想快速复现非配对图像转换论文、做基线对比的实验型开发者另一类是准备在真实业务里做风格迁移需要自己动手换网络结构、调损失权重的工程型开发者。这个项目把论文里的公式变成了可读的 Python 代码跑通它你就理解了 CycleGAN 的每一个关键决策。2. 技术栈与代码结构一个 PyTorch 工程模板的完整拆解2.1 依赖清单这套代码踩过的坑项目的技术栈非常轻核心依赖就是 PyTorch 和 torchvision再加 Pillow、numpy 这些图像处理必备库。没有引入任何重型框架也没有用 HuggingFace 之类的抽象层所有模型、损失、训练循环都写在原生 PyTorch 里这样二次开发时不会被框架约束。我一般用这个版本的依赖组合python3.8 torch1.9.0 torchvision0.10.0 numpy Pillow tqdm tensorboard2.5.0为什么要求这么低因为 CycleGAN 本身的网络结构不挑卡显存吃紧的时候甚至可以用 CPU 跑通一个小轮次只是速度感人。Pytorch 1.9 以上的InstanceNorm2d接口稳定后续版本没有破坏性变更所以这个下限足够。如果你用的是最新版 PyTorch直接装最新的 torchvision 也没问题代码里没有用到会被废弃的 API。tensorboard 不是训练必需但建议装上后面讲可视化时会用到。如果你不想装把--use_tensorboard参数关掉即可代码里做了兼容处理。2.2 目录结构每个文件负责什么拿到源码包后第一件事就是看目录。这个项目没有把什么train.py、test.py全堆在根目录而是按模块划分我非常喜欢这种结构CycleGAN/ ├── main.py # 入口脚本解析参数并启动训练/测试 ├── models/ │ ├── __init__.py # 模型统一导入 │ ├── generator.py # 生成器ResNet9 残差块 │ └── discriminator.py # 判别器PatchGAN ├── data/ │ ├── __init__.py │ └── dataset.py # 自定义 Dataset读取两个域 ├── utils/ │ ├── __init__.py │ ├── loss.py # 损失函数封装 │ ├── visualization.py # TensorBoard 图像记录 │ └── utils.py # checkpoint 保存/加载、评分 ├── checkpoints/ # 模型权重保存目录 └── datasets/ ├── trainA/ # A 域训练图片 ├── trainB/ # B 域训练图片 ├── testA/ # A 域测试图片 └── testB/ # B 域测试图片这个结构对应了 CycleGAN 的两大核心模块生成器和判别器分开文件训练入口只做组合数据加载单独放这样任何时候替换数据集都不需要动模型代码。checkpoints目录会自动生成不需要手动建main.py启动时会检查。utils/loss.py值得单独提一下它把 GAN 损失、循环损失、身份损失都封装成函数避免了在训练循环里重复写同样的 MSE 和 L1 公式。你在改实验时直接改这一个文件就够了不用绕到训练循环里到处找。2.3 数据加载不用 ImageFolder自己写 Dataset原版的 CycleGAN 数据加载有一个关键点图像不需要配对但需要一个 Dataset 能够同时返回 A 域和 B 域的图像。torchvision.datasets.ImageFolder只能加载一个域而且会自动按文件夹分类不适合这里的使用方式。所以项目里的data/dataset.py自己继承了torch.utils.data.Dataset。看关键代码# data/dataset.py import os from PIL import Image from torch.utils.data import Dataset from torchvision import transforms class CycleGANDataset(Dataset): def __init__(self, root_a, root_b, img_size256): self.files_a sorted(os.listdir(root_a)) self.files_b sorted(os.listdir(root_b)) self.root_a root_a self.root_b root_b self.transform transforms.Compose([ transforms.Resize((img_size, img_size)), transforms.ToTensor(), transforms.Normalize(mean[0.5, 0.5, 0.5], std[0.5, 0.5, 0.5]) ]) def __len__(self): return max(len(self.files_a), len(self.files_b)) def __getitem__(self, index): name_a self.files_a[index % len(self.files_a)] name_b self.files_b[index % len(self.files_b)] img_a Image.open(os.path.join(self.root_a, name_a)).convert(RGB) img_b Image.open(os.path.join(self.root_b, name_b)).convert(RGB) return self.transform(img_a), self.transform(img_b)逻辑说明__len__返回两个目录中较大的那个长度配合取模操作让数量少的域也能被完整遍历到不会因为其中一个域只有 100 张图而让另一个域只被用 100 次。__getitem__每次按当前 index 分别从两个域各取一张图注意这里用的是同一个 index 加上各自长度的取模所以 A 域和 B 域的样本在每次迭代时是随机组合的完全不需要配对关系。参数说明root_a和root_b是 A、B 两个域的图片文件夹路径img_size默认 256这是 CycleGAN 论文中的标准分辨率。归一化用了mean0.5, std0.5会把像素从[0,1]映射到[-1,1]这和生成器输出层的Tanh激活函数是对应的。使用Image.open后调用.convert(RGB)是为了去掉 PNG 的透明通道也可以避免灰度图混入时导致通道数不一致。如果你需要做数据增强可以在这个transform里插入RandomHorizontalFlip但原版 CycleGAN 只在训练时对单张图做水平翻转这个项目把增强参数也开放给了命令行后面会看到。3. 核心模型搭建生成器、判别器和损失函数的完整实现3.1 生成器ResNet9 残差块与反射填充CycleGAN 的生成器是一个编码器-转换器-解码器结构。输入是 3 通道的 256x256 图像先用两层卷积下采样到 64x64然后经过 9 个残差块进行域转换再用两层转置卷积上采样回 256x256。这里最关键的是残差块内部它用了ReflectionPad2d而不是普通零填充这是图像风格迁移的一个经典血泪经验。看models/generator.py里的残差块实现# models/generator.py import torch.nn as nn class ResidualBlock(nn.Module): def __init__(self, in_channels): super().__init__() self.block nn.Sequential( nn.ReflectionPad2d(1), nn.Conv2d(in_channels, in_channels, kernel_size3, biasFalse), nn.InstanceNorm2d(in_channels), nn.ReLU(inplaceTrue), nn.ReflectionPad2d(1), nn.Conv2d(in_channels, in_channels, kernel_size3, biasFalse), nn.InstanceNorm2d(in_channels), ) def forward(self, x): return x self.block(x)逻辑说明残差块内部是两个卷积每个卷积前垫一层反射填充。ReflectionPad2d(1)会把边界的像素镜像反射到外面比如图像最左一列像素会出现在新矩阵的最左外侧。相比ZeroPad2d反射填充不会在边缘引入黑色伪影这对图像转换任务很重要。第二个卷积层的输出self.block(x)和输入x相加形成恒等映射路径。残差块不改变通道数和空间尺寸所以可以串联多个而不出维度问题。参数说明in_channels是残差块的通道数在标准 CycleGAN 中编码器下采样后通道数变成 256所以 ResNet 块处理的就是 256 通道。卷积层设置biasFalse因为它后面紧跟InstanceNorm2d归一化层会做平移操作偏置会被屏蔽留着反而浪费参数。使用InstanceNorm2d而不是BatchNorm2d是因为 CycleGAN 通常使用较小的 batch size甚至经常是 batch size1BN 在单样本时统计均值方差没有意义而 InstanceNorm 对单张图归一化效果更稳定。生成器的主体就是用一个nn.ModuleList把这 9 个残差块串起来然后在forward里依次调用。注意生成器输入的值域是[-1,1]输出层用了nn.Tanh()把输出也压到[-1,1]和判别器的输入范围保持一致。3.2 判别器70x70 PatchGAN 的秘密判别器不是输出一个 0 到 1 的标量而是输出一个 NxN 的特征图每个点对应输入图像上一个 70x70 的感受野区域。这个设计叫 PatchGAN它的优势是参数少、更关注局部纹理同时能逼着生成器把每个局部区域都做得逼真而不是只求整体分布像。看models/discriminator.py代码# models/discriminator.py import torch.nn as nn class PatchDiscriminator(nn.Module): def __init__(self, in_channels3): super().__init__() self.model nn.Sequential( nn.Conv2d(in_channels, 64, kernel_size4, stride2, padding1), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv2d(64, 128, kernel_size4, stride2, padding1), nn.InstanceNorm2d(128), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv2d(128, 256, kernel_size4, stride2, padding1), nn.InstanceNorm2d(256), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv2d(256, 1, kernel_size4, stride1, padding1), ) def forward(self, x): return self.model(x)逻辑说明输入是 256x256 的 3 通道图像经过三组 stride2 的卷积后空间尺寸逐层减半256 - 128 - 64 - 32。最后一层卷积的 stride 改为 1输出通道数为 1最终得到一个 31x31 的特征图。这个特征图上的每个像素都对应原图的一个局部区域判别器训练的目标是让所有真实图像对应的像素接近 1生成的图像对应的像素接近 0。因为判别器只看局部所以它能指导生成器提高高频纹理的真实度。参数说明卷积核统一用kernel_size4这是 DCGAN 之后常用的配置配合padding1能精确控制尺寸变化。LeakyReLU(0.2)的负斜率是 0.2避免 ReLU 把负值全部置零导致梯度消失。每个卷积后都接InstanceNorm2d只有第一层不加这是稳定训练的一个常见约定——对原始输入做归一化容易破坏数据分布。倒数第二层和最后一层之间没有归一化直接输出 logits这个 logits 会传给 MSE 损失。这里有个易错点判别器的输出不是经过 Sigmoid 的概率而是裸的 logits所以后面计算损失时用的MSE_loss直接拿输出和目标值比不能再过 Sigmoid否则梯度会不稳定很多新手在这里翻车。3.3 损失函数三驾马车的平衡CycleGAN 的损失由三部分组成对抗损失、循环一致性损失、身份损失。对抗损失用 LSGAN 的 MSE 形式循环一致性损失用 L1身份损失也用 L1。看utils/loss.py的封装# utils/loss.py import torch import torch.nn.functional as F def gan_loss(real_pred, fake_pred, is_discriminatorTrue): 计算 LSGAN 损失 if is_discriminator: real_loss F.mse_loss(real_pred, torch.ones_like(real_pred)) fake_loss F.mse_loss(fake_pred, torch.zeros_like(fake_pred)) return (real_loss fake_loss) * 0.5 else: return F.mse_loss(fake_pred, torch.ones_like(fake_pred)) def cycle_loss(rec_a, real_a, rec_b, real_b): return F.l1_loss(rec_a, real_a) F.l1_loss(rec_b, real_b) def identity_loss(fake_b, real_b, fake_a, real_a): return F.l1_loss(fake_b, real_b) F.l1_loss(fake_a, real_a)逻辑说明gan_loss在判别器模式下把真实图像判别器输出尽可能逼近 1把生成图像输出尽可能逼近 0并取平均在生成器模式下只计算生成图像输出逼近 1 的损失即生成器想让判别器错判。cycle_loss计算 A 域图经过 G 和 F 再回到 A 域的差异以及 B 域图经过 F 和 G 再回到 B 域的差异。L1 损失比 L2 更尖锐能保留更多边缘信息所以在循环一致性上选 L1。identity_loss是可选的它的作用是让生成器在输入已经属于目标域时尽量别改比如输入一张冬天的图经过冬天到夏天再回来的网络时应该尽量保持原样。这个损失权重要调小否则生成器太平庸。为什么循环一致性损失能解决非配对问题因为两个生成器互为逆映射。G 把 A 变 BF 把 B 变 A如果 G(A) 再经过 F 能还原出 A就说明 G 没有随意发挥它学会的映射是受约束的。这就是 CycleGAN 和单纯 GAN 的本质区别。在main.py中这些损失按权重组合成总生成器损失loss_G loss_gan_AB loss_gan_BA 10.0 * loss_cycle 0.5 * loss_identity权重 10 和 0.5 是论文里的默认值项目也做成了命令行参数。这个权重对训练稳定性影响极大后面避坑里会单独展开。4. 训练流程与关键参数从命令行开始跑通4.1 入口 main.py从参数到模型初始化项目的入口是main.py它用标准库argparse解析参数。整个训练脚本只需要一行命令就能启动所有可调项全暴露出来。python main.py --mode train --root_a ./datasets/trainA --root_b ./datasets/trainB --epochs 200 --batch_size 1 --image_size 256这句命令的含义指定 A 域训练图片在datasets/trainA目录B 域在datasets/trainB训练 200 个 epochbatch size 为 1图片缩放到 256x256。如果你显存不够可以把image_size调成 128训练速度会快很多但生成图像细节会有所损失。模型初始化部分main.py会根据参数分别构建两个生成器和两个判别器# main.py 初始化片段 from models.generator import ResnetGenerator from models.discriminator import PatchDiscriminator gen_ab ResnetGenerator(num_resblocks9) gen_ba ResnetGenerator(num_resblocks9) dis_a PatchDiscriminator(in_channels3) dis_b PatchDiscriminator(in_channels3)ResnetGenerator内部会自动判断num_resblocks9 个块是论文原版配置如果改成 6 个可以加快训练但效果略降。两个生成器结构完全一样但权值不相干两个判别器也完全一样。这里有一个容易忽略的设计gen_ab负责 A 到 Bgen_ba负责 B 到 A命名上严格区分后面对称更新时就不会晕。初始化之后代码会检查checkpoints目录下是否存在已有权重如果有就会自动加载实现断点续训。4.2 训练循环谁先更新谁训练循环是整个项目最值得细看的部分因为它决定了两套网络能否稳定博弈。CycleGAN 里生成器和判别器不是同时更新而是每轮先更新生成器再更新判别器并且判别器更新时使用detach()切断生成器的梯度。看核心代码# main.py 训练循环简化版 for epoch in range(start_epoch, epochs): for i, (real_a, real_b) in enumerate(train_loader): # 第一步更新生成器 fake_b gen_ab(real_a) fake_a gen_ba(real_b) rec_a gen_ba(fake_b) rec_b gen_ab(fake_a) id_b gen_ab(real_b) id_a gen_ba(real_a) loss_gan_ab gan_loss(dis_b(real_b), dis_b(fake_b), is_discriminatorFalse) loss_gan_ba gan_loss(dis_a(real_a), dis_a(fake_a), is_discriminatorFalse) loss_cycle 10.0 * cycle_loss(rec_a, real_a, rec_b, real_b) loss_id 0.5 * identity_loss(id_b, real_b, id_a, real_a) loss_G loss_gan_ab loss_gan_ba loss_cycle loss_id gen_optimizer.zero_grad() loss_G.backward() gen_optimizer.step() # 第二步更新判别器 fake_b_det fake_b.detach() fake_a_det fake_a.detach() loss_dis_a gan_loss(dis_a(real_a), dis_a(fake_a_det), is_discriminatorTrue) loss_dis_b gan_loss(dis_b(real_b), dis_b(fake_b_det), is_discriminatorTrue) loss_D (loss_dis_a loss_dis_b) * 0.5 dis_optimizer.zero_grad() loss_D.backward() dis_optimizer.step()逻辑说明生成器更新时输入真实图像real_a和real_b前向得到fake_b、fake_a再各自送入反向生成器得到重构图像rec_a、rec_b。循环一致性损失要求rec_a逼近real_arec_b逼近real_b。身份损失要求gen_ab(real_b)逼近real_bgen_ba(real_a)逼近real_a。这代代码一次性算了生成器的所有损失然后反向传播更新两个生成器的共享优化器。到判别器更新时fake_b和fake_a已经由生成器计算出来了但必须.detach()让梯度不会流向生成器。如果不 detach判别器更新时就会顺带更新生成器参数导致两把刀在同一个回合里互相砍训练很容易炸掉。这也是一个常见翻车点。参数说明gen_optimizer和dis_optimizer都是 Adam学习率默认 2e-4beta10.5。GAN 类任务常用 beta10.5 而不是默认的 0.9这样能减小动量影响让判别器更快响应生成器的变化。优化器在初始化后会对所有参数管理不需要单独给每个网络分配一个优化器两个生成器共用一个、两个判别器共用一个能减小收敛震荡。4.3 学习率调度与 checkpoint 保存训练到一半时图像质量往往还不理想。CycleGAN 原版采用线性衰减学习率的策略前 100 个 epoch 保持 2e-4后 100 个 epoch 线性降到 0。这个动作通过torch.optim.lr_scheduler.LambdaLR实现# main.py 调度器配置 def lambda_rule(epoch): num_epochs 200 decay_start 100 return 1.0 if epoch decay_start else 1.0 - (epoch - decay_start) / (num_epochs - decay_start) scheduler_gen torch.optim.lr_scheduler.LambdaLR(gen_optimizer, lr_lambdalambda_rule) scheduler_dis torch.optim.lr_scheduler.LambdaLR(dis_optimizer, lr_lambdalambda_rule)epoch从 0 开始计数所以前 100 个 epoch 学习率不变从第 100 个 epoch 开始线性下降到第 199 个 epoch 正好降成 0。这样做的好处是后期学习率小能让生成器在稳定的判别器下精细调整纹理而不是大步跨越把权重踢到坏区域。checkpoint 保存也写在训练循环里每 5 个 epoch 自动保存一次到checkpoints目录。保存的字典里同时包含两个生成器、两个判别器的state_dict还有优化器状态和当前 epoch 数这样恢复训练时不会丢失动量参数# main.py checkpoint 保存片段 torch.save({ gen_ab: gen_ab.state_dict(), gen_ba: gen_ba.state_dict(), dis_a: dis_a.state_dict(), dis_b: dis_b.state_dict(), gen_optimizer: gen_optimizer.state_dict(), dis_optimizer: dis_optimizer.state_dict(), epoch: epoch 1, }, checkpoints/latest.pt)加载时只需要把torch.load返回的字典逐项load_state_dict然后把epoch作为起始 epoch 继续循环。注意 PyTorch 默认在 GPU 环境下保存的张量会带着设备信息如果你换卡训练加载时要加map_locationcpu再转移到新设备否则容易报错。5. 避坑指南CycleGAN 训练中的常见问题排查5.1 训练不收敛损失振荡严重现象训练几十个 epoch 后生成器损失和判别器损失来回横跳生成的图像忽好忽坏甚至出现噪点。原因最常见的是学习率太高或者判别器更新步数比生成器少导致判别器跟不上生成器的变化节奏。另一个原因是输入图像的归一化不一致比如只对 A 域做了标准化B 域还是[0,1]范围判别器会直接把亮度当成区分信号导致训练震荡。解决先把学习率降到 1e-4同时确认两个域都用Normalize(mean0.5, std0.5)让输入落在[-1,1]。然后观察每个 step 里两个损失的数量级如果判别器损失长期低于生成器损失说明判别器太强可以降低判别器相关权重或把判别器更新次数从 1 改成每 2 个 step 更新一次项目支持--n_critic参数。5.2 生成图像发糊纹理细节全部丢失现象生成的图像整体布局正确但表面像蒙了一层雾边缘模糊。原因循环一致性损失权重太高模型为了满足还原重构而选择保守的映射输出趋近于均值。另外生成器最后一层用nn.Tanh()输出如果你在损失计算时又对输出做了一次sigmoid或归一化会让梯度变得非常小生成器学不动细节。解决把--lambda_cycle从 10 降到 5适当降低重构约束同时检查损失函数代码确认生成器输出直接参与 MSE没有额外激活。我一般还会把判别器的感受野调大一点用--patch_size参数控制比如从 70 改成 1即判别器只看整体这会牺牲局部纹理但能恢复宏观结构。如果模糊严重优先怀疑身份损失权重把它从 0.5 降到 0.2。5.3 生成器“罢工”所有输出都变成同一张图现象训练到中后期无论输入什么图像生成器输出几乎相同而且循环一致性损失一直卡住不降。原因这是 GAN 训练里经典的 mode collapse生成器发现输出一个“万能图”就能骗过判别器就不再学习输入变化。根本原因是判别器能力太弱或者对抗损失占主导生成器走捷径。解决先停掉判别器的 InstanceNorm或者把判别器的通道数减半让判别器变弱迫使生成器必须提供更丰富的信息来应对。同时把循环一致性权重调大它是对抗损失的“正则项”能有效阻止生成器忽略输入。项目里通过--lambda_identity也能略加帮助但我发现最有效的是把 batch size 从 1 提到 4让生成器在同一个 batch 内看到多个不同样本天然不会坍缩到单点。5.4 显存溢出训练直接中断现象启动训练十几秒后程序报CUDA out of memory。原因256x256 的输入经过两个生成器和两个判别器中间还要保存多份特征图用于反向传播显存占用轻松超过 8GB。分辨率越高通道数越多开销呈平方级上涨。解决第一选择是把--image_size降到 128这在很多风格迁移任务里已经足够训练结束再用 256 微调。第二选择是打开--grad_accumulation功能通过梯度累积模拟更大的 batch 来减少显存峰值。项目在训练循环里实现了 accumulate 变量每攒到--accumulation_steps次才做一次优化器 step。如果还不行把判别器的通道数从 64 改成 32显存占用接近减半代价是判别力弱一点。5.5 数据量太少过拟合严重现象训练集每个域只有几十张图训练后期生成器把训练图像原样背下来测试集上效果一塌糊涂。原因CycleGAN 是很吃数据的模型原版论文每个域用了上千张图片。数据少时生成器太容易记住有限样本尤其是循环一致性损失会加剧这种行为。解决项目内置了随机水平翻转和随机裁剪增强打开--preprocessflip_crop会让每张图每次迭代都被随机裁剪出一个 256x256 区域相当于数据量放大几倍。如果还不行就降低生成器残差块数量比如从 9 改成 6减少模型容量同时把--lambda_cycle降到 5避免过强的重构约束逼着模型死记硬背。这种问题没有完美解最靠谱的还是多去收集图像哪怕不配对都比全靠增强硬撑强。6. 二次开发技巧换数据集、调损失、加可视化6.1 把自己的业务图替换进去假设你现在要把工程图纸里的手绘草图转换成彩色渲染图。只需要准备两个文件夹A放草图B放渲染图。图片不必一一对应但要做到内容域一致比如所有草图和渲染图都应该显示同一个种类的椅子否则模型学不到稳定的映射关系。目录组织好后直接执行训练命令如果发现生成图偏色先调--lambda_identity如果发现结构不对调--lambda_cycle。这是最快见效的改造。6.2 生成器换成 UNet 或减少残差块如果任务很接近配对的图像转换或者两个域的差异只是颜色、纹理9 个残差块有点重。在models/generator.py里加一个UnetGenerator类编码器用四层下采样解码器对应四层上采样中间用跳跃连接。修改时注意UNet 的输入和输出尺寸要一致跳跃连接需要把编码器层和下采样层配对保存。改完之后在main.py的初始化处加一个--generator_type参数默认resnet传unet时走新分支。如果你不想动文件把残差块数量改成 6 也能加速 30% 左右很多颜色偏移任务用 6 块就够了。6.3 用 TensorBoard 盯住中间产物训练不看中间图像等于闭着眼开车。项目在utils/visualization.py里已经写好了一个log_images函数它会把real_a、fake_b、rec_a、real_b、fake_a、rec_b六张图拼成网格写到 TensorBoard# utils/visualization.py from torchvision.utils import make_grid from torch.utils.tensorboard import SummaryWriter def log_images(writer, tag, real_a, fake_b, rec_a, real_b, fake_a, rec_b, step): images { real_a: real_a, fake_b: fake_b, rec_a: rec_a, real_b: real_b, fake_a: fake_a, rec_b: rec_b, } for name, im in images.items(): writer.add_image(f{tag}/{name}, make_grid(im.clamp(-1, 1).detach().cpu(), nrow4), step)调试时我还会额外保存gen_ab(real_b)和gen_ba(real_a)的身份输出这样能直观看到身份损失到底有没有起作用。如果发现fake_b和real_a看起来几乎一样而rec_a严重偏差说明循环一致性权重过大如果fake_b清晰度高但内容完全不对则要加循环权重。这些都可以靠 TensorBoard 曲线来定位。从那以后我每次换数据集训练都会强制走一遍同样的流程先跑 5 个 epoch 看图像变化幅度确认生成器在动再跑 50 个 epoch 观察损失曲线最后再根据可视化微调权重。这套习惯帮我避开了好几轮“损失看着在降图像其实不学”的假收敛。希望帮到你。本文还有配套的精品资源点击获取
阅读完成 · 觉得有帮助？
