☰
pix2pix原理与PyTorch实现:U-Net生成器和PatchGAN判别器实战
2026/10/18 9:58:19 网站建设 项目流程

开门见山说一句:pix2pix是我在GAN这条路上第一个真正“用明白”的模型。如果你看过一些图像生成项目,八成见过那种“线稿一键上色”“街区地图转卫星图”“白天照片转夜景”的效果演示,背后基本都是它。pix2pix属于条件生成对抗网络(Conditional GAN)里最经典的实现之一,输入一张图,输出一张图,端到端搞定,不需要手工设计特征。这篇博文我会从原理出发,带你把整个模型在PyTorch里从零写出来,包括网络结构、损失函数、训练循环,以及我在实际训练里踩过的那些坑。

1. pix2pix在做什么:图像翻译问题与设计思路

1.1 从一个朴素的直觉说起:什么是图像到图像的翻译

先把问题定义说清楚。图像翻译(Image-to-Image Translation)指的是:给定一张输入图像,希望模型生成一张与之对应的输出图像。比如你有一张边缘线稿,想得到一张完整的上色照片;你有一张灰度图,想得到一张彩色图;你有一张日间街景,想得到一张夜间街景。输入输出的像素之间存在某种空间上的一一对应关系,这跟“根据一段文字生成一张图”完全不同,后者没有严格的像素对齐约束。

pix2pix把这个任务统一成一个框架,模型看到的只有输入图像,输出的就是目标域图像。它不关心你具体是做什么翻译任务,只要你准备成“输入-目标”成对的数据,它就能训练。这个通用性正是它的魅力所在。当年这篇工作把一堆看起来毫无关联的任务都纳入了同一个模型体系,语义分割图转街景照片、深度图转真实场景、灰图上色,全都用同一套网络结构和训练流程跑通。

1.2 为什么普通CNN不够用:结构和语义都要对齐

如果只是想把一张输入图映射成一张输出图,最自然的想法是用一个CNN做回归。输入边缘图,输出彩色图,像素级算一个L2 Loss,网络自己学映射。但纯CNN做这类任务有个问题:它倾向于输出“平均化”的结果,遇到一个输入可能对应多种合理输出的情况时,模型会选择它们的平均值。

举个好理解的例子:一条边线勾勒的轮廓,可能是猫也可能是狗,纯回归模型在模糊地带会生成一张“既像猫又像狗”的模糊图。因为L2/L1损失函数鼓励模型在所有可能答案里取折中,折中出来的图像就是糊的、纹理丢失的。真实图像生成需要的是纹理和细节,不是“像素平均”。这也是为什么pix2pix没有只使用像素损失,而是引入判别器去逼迫生成器输出以假乱真的高频细节。

1.3 引入对抗损失:生成器骗过判别器

GAN的思路是训练一个生成器G和一个判别器D。D的任务是区分输入图像是“真实目标图像”还是“G生成出来的假图像”,G的任务则是让生成图像足够真实,骗过D。两者不断博弈,G的输出会越来越好。

pix2pix用的是条件GAN,也就是把输入图像x同时喂给G和D。生成器G的输入是x,输出是假图像G(x);判别器D的输入是“x拼接上真实目标y”或者“x拼接上假图像G(x)”,输出一个判断分数,表示这组配对有多真。这样一来,D不仅要判断图像像不像真的,还要判断它是否跟输入x匹配。这个“条件”约束非常关键,否则D只学会判断“照片是不是真实照片”,无视输入,生成结果就不可控了。

1.4 网络结构选型的核心逻辑:U-Net加PatchGAN

确定了GAN框架之后,剩下两个关键选型:生成器用什么骨架,判别器怎么设计。pix2pix的答案是:生成器用U-Net结构,判别器用PatchGAN。

U-Net的编码器-解码器结构自带跨层跳连,可以把浅层的高频细节(边缘、纹理)直接传给解码器,解决图像翻译里“输入输出需要空间对齐”的痛点。PatchGAN则不做整图判断,而是对图像的一个个局部块做真假分类。这两点我会在第3章详细展开代码实现。先记住一句话:U-Net管“信息传递”,PatchGAN管“局部纹理真实感”,两者各司其职。

2. 环境准备与关键组件选型

2.1 最小可跑环境清单

在动手写代码前,先把环境搞定。我用的配置是Python 3.8、PyTorch 1.10以上、CUDA 11.x,一块8GB显存的显卡就能跑小尺寸的pix2pix训练。如果显存不够,可以把输入切到256×256以下,或者减小batch size。

依赖库方面,核心是torch、torchvision、numpy、pillow、matplotlib。数据加载和可视化用这几个就够了,不需要额外装复杂的深度学习框架。建议你用conda单独建一个环境,避免把系统Python搞乱。

conda create -n pix2pix python=3.8 conda activate pix2pix pip install torch torchvision numpy pillow matplotlib

实测下来,这套组合在Windows和Linux上都能稳定跑。需要提醒的是,PyTorch版本太老(比如1.4)会导致一些API不兼容,特别是后面写自定义Dataset时,保持较新版本会省去很多调试时间。

2.2 数据集选择与预处理

pix2pix需要成对数据,也就是“输入图像”和“目标图像”一一对应。官方那批数据集因为文件结构清晰、任务多样,经常被当成入门首选。你按任务选即可:

任务输入输出适合入门程度
建筑立面生成语义标签图真实建筑照片入门友好
街景翻译街景分割图真实街景经典任务
灰度上色灰度图彩色图数据易构造
线稿上色边缘线稿彩色照片效果直观

预处理做的第一件事是统一尺寸。pix2pix官方实现里常见做法是先将图片缩放到286×286,再随机裁剪成256×256。为什么先放大再裁剪?这相当于一种数据增强,能增加平移不变性,让模型对目标的位移鲁棒一点。如果你偷懒只做resize到256×256,训练时模型见过的样本多样性会明显下降,测试时的泛化能力也跟着打折。

归一化方面,把输入和目标的像素值都从[0,255]缩放到[0,1],再减0.5除0.5变成[-1,1]。这个区间跟生成器最后一层Tanh的输出范围一致,可以保证生成图像和真实图像处于同一个数值尺度。很多人在这块容易踩坑:如果忘了把真实目标图像同步做归一化,判别器会瞬间拉开真假差异,训练直接崩。

2.3 自定义Dataset与数据加载器实现

PyTorch里自定义数据集最常见的写法是继承torch.utils.data.Dataset。我在实际项目里习惯把输入和目标的文件夹分开存放,两张同名图片配对。这样做的好处是数据管理直观,换任务时不用改数据接口。

import os from PIL import Image from torch.utils.data import Dataset import torchvision.transforms as T class PairedDataset(Dataset): def __init__(self, input_dir, target_dir, size=256): self.input_dir = input_dir self.target_dir = target_dir self.filenames = sorted(os.listdir(input_dir)) self.size = size def __len__(self): return len(self.filenames) def __getitem__(self, idx): name = self.filenames[idx] input_path = os.path.join(self.input_dir, name) target_path = os.path.join(self.target_dir, name) input_img = Image.open(input_path).convert('RGB') target_img = Image.open(target_path).convert('RGB') # 先统一缩放到286,再随机裁剪到256,等价于数据增强 resize = T.Resize((286, 286), interpolation=T.InterpolationMode.BICUBIC) input_img = resize(input_img) target_img = resize(target_img) seed = torch.randint(0, 10000, (1,)).item() torch.manual_seed(seed) crop = T.RandomCrop((self.size, self.size)) input_img = crop(input_img) torch.manual_seed(seed) target_img = crop(target_img) # 随机水平翻转,同样需要保持同步 if torch.rand(1).item() > 0.5: input_img = T.functional.hflip(input_img) target_img = T.functional.hflip(target_img) # 归一化到[-1,1] to_tensor = T.ToTensor() norm = T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) input_img = norm(to_tensor(input_img)) target_img = norm(to_tensor(target_img)) return input_img, target_img

这里有个细节我得专门强调:random seed的同步。因为同一组输入和目标必须经过完全相同的裁剪和翻转,不能输入裁完之后目标裁的位置不一样,那样空间对齐就废了。除了用torch.manual_seed同步,你也可以用torchvision.transforms里的RandomCrop配合concat图片再切回去的思路,本质一样。我推荐seed同步法,代码最简单。

写好了Dataset之后,DataLoader的配置也有讲究。batch size一般取1到4,pix2pix因为显存占用大,小batch很常见。注意num_workers根据机器CPU核数设置,Windows下如果多进程加载报错,先归零跑通再逐步加大。

3. 生成器与判别器的PyTorch实现

3.1 U-Net生成器:跳连如何保住细节

pix2pix的生成器U-Net由编码器(encoder)和解码器(decoder)构成。encoder不断做卷积和下采样,把图像压缩成高维特征;decoder把特征逐级上采样回原分辨率。关键在跳连(skip connection):encoder第i层输出的feature map,直接拼接到decoder对应层级上。

为什么必须有跳连?因为单独靠瓶颈层的特征做上采样,只能恢复语义,恢复不了高频细节。边缘、墙角线、纹理这些信息在浅层最丰富,经过多层下采样后基本丢失了。跳连相当于给解码器开了一条直达通道,把浅层细节带到深层。有人把U-Net类比为“建筑中从一楼到五楼的直达电梯”,虽然信息从楼梯逐层压缩了,但关键楼层之间还有快速通路,信息不会丢。

import torch import torch.nn as nn class UNetBlock(nn.Module): def __init__(self, in_ch, out_ch, use_dropout=False): super().__init__() self.block = nn.Sequential( nn.Conv2d(in_ch, out_ch, kernel_size=4, stride=2, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), ) self.use_dropout = use_dropout if use_dropout: self.dropout = nn.Dropout(0.5) def forward(self, x): x = self.block(x) if self.use_dropout: x = self.dropout(x) return x

等等,这段代码是为了展示卷积块的通用写法,但实际U-Net的encoder和decoder行为不一样。为了完整性,下面给你看完整的U-Net实现,注释里会标清楚每一段的作用。

class UnetGenerator(nn.Module): def __init__(self, in_ch=3, out_ch=3, num_filters=64): super().__init__() # encoder: 输入 -> 逐级下采样 self.e1 = self._encoder_block(in_ch, num_filters, batchnorm=False) # 256->128 self.e2 = self._encoder_block(num_filters, num_filters*2, batchnorm=True) # 128->64 self.e3 = self._encoder_block(num_filters*2, num_filters*4, batchnorm=True) # 64->32 self.e4 = self._encoder_block(num_filters*4, num_filters*8, batchnorm=True) # 32->16 self.e5 = self._encoder_block(num_filters*8, num_filters*8, batchnorm=True) # 16->8 self.e6 = self._encoder_block(num_filters*8, num_filters*8, batchnorm=True) # 8->4 self.e7 = self._encoder_block(num_filters*8, num_filters*8, batchnorm=True) # 4->2 self.e8 = self._encoder_block(num_filters*8, num_filters*8, batchnorm=True) # 2->1 # decoder: 瓶颈 -> 逐级上采样 self.d1 = self._decoder_block(num_filters*8, num_filters*8, dropout=True) # 1->2 self.d2 = self._decoder_block(num_filters*16, num_filters*8, dropout=True) # 2->4 self.d3 = self._decoder_block(num_filters*16, num_filters*8, dropout=True) # 4->8 self.d4 = self._decoder_block(num_filters*16, num_filters*8, dropout=False) # 8->16 self.d5 = self._decoder_block(num_filters*16, num_filters*4, dropout=False) # 16->32 self.d6 = self._decoder_block(num_filters*8, num_filters*2, dropout=False) # 32->64 self.d7 = self._decoder_block(num_filters*4, num_filters, dropout=False) # 64->128 self.d8 = self._decoder_block(num_filters*2, num_filters, dropout=False, last=True) # 128->256 self.tanh = nn.Tanh() def _encoder_block(self, in_ch, out_ch, batchnorm=True): layers = [nn.Conv2d(in_ch, out_ch, kernel_size=4, stride=2, padding=1, padding_mode='reflect', bias=False)] if batchnorm: layers.append(nn.BatchNorm2d(out_ch)) layers.append(nn.LeakyReLU(0.2, inplace=True)) return nn.Sequential(*layers) def _decoder_block(self, in_ch, out_ch, dropout=False, last=False): layers = [nn.ConvTranspose2d(in_ch, out_ch, kernel_size=4, stride=2, padding=1, bias=False)] if last: layers.append(nn.Tanh()) else: layers.append(nn.BatchNorm2d(out_ch)) layers.append(nn.ReLU(inplace=True)) if dropout: layers.append(nn.Dropout(0.5)) return nn.Sequential(*layers) def forward(self, x): # encoder e1 = self.e1(x) e2 = self.e2(e1) e3 = self.e3(e2) e4 = self.e4(e3) e5 = self.e5(e4) e6 = self.e6(e5) e7 = self.e7(e6) e8 = self.e8(e7) # decoder,注意cat的顺序 d1 = self.d1(e8) d1 = torch.cat([d1, e7], dim=1) d2 = self.d2(d1) d2 = torch.cat([d2, e6], dim=1) d3 = self.d3(d2) d3 = torch.cat([d3, e5], dim=1) d4 = self.d4(d3) d4 = torch.cat([d4, e4], dim=1) d5 = self.d5(d4) d5 = torch.cat([d5, e3], dim=1) d6 = self.d6(d5) d6 = torch.cat([d6, e2], dim=1) d7 = self.d7(d6) d7 = torch.cat([d7, e1], dim=1) d8 = self.d8(d7) return self.tanh(d8)

从U-Net的实现可以看出,encoder用的是LeakyReLU,decoder用的是ReLU。LeakyReLU在编码阶段能保留负半轴信息,加快收敛;解码阶段输出需要非负特征,ReLU更合适。最后用Tanh把输出限制到[-1,1],正好跟数据归一化范围对上。

有一个容易写错的地方:通道数。跳连concat后通道数是当前层输出加对应encoder层输出,比如decoder第二层输入应该是8*nf+8*nf=16*nf。我见过不少初学者在这里想当然地写成8*nf,结果维度对不上报错。上面代码里的通道数我都仔细核对过了,照着抄能跑。

3.2 PatchGAN判别器:为什么局部判断就够了

如果说生成器管“怎么生成”,判别器管“怎么评价”。pix2pix判别器名叫PatchGAN,分类器不对整张图输出一个真/假分数,而是输出一个N×N的矩阵,矩阵上的每个元素对应原始图像的某个感受野区域。

这个设计很多人第一反应是困惑:为什么不直接全图判断?深层逻辑是,图像翻译里最重要的其实是“局部纹理是否逼真”,而全局语义是否一致由L1损失负责。PatchGAN只用把精力集中在一个小区域,判别的难度降低,梯度信息也更细粒度。整图判断的时候,只要图片整体看着像,判别器就给高分,生成器就能在一些小区域蒙混过关。PatchGAN把图像分成一片片来处理,相当于考试时一道题一道题批改,而不是整张卷子给个总分。

PatchGAN的感受野大小取决于网络层数。典型实现是三层卷积加一层输出卷积,每层stride=2,感受野大约70×70。输入一张256×256的图,输出是30×30的矩阵,前文中说N×N即如此。

class PatchDiscriminator(nn.Module): def __init__(self, in_ch=3, num_filters=64): super().__init__() self.model = nn.Sequential( nn.Conv2d(in_ch*2, num_filters, kernel_size=4, stride=2, padding=1, padding_mode='reflect'), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(num_filters, num_filters*2, kernel_size=4, stride=2, padding=1, padding_mode='reflect', bias=False), nn.BatchNorm2d(num_filters*2), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(num_filters*2, num_filters*4, kernel_size=4, stride=2, padding=1, padding_mode='reflect', bias=False), nn.BatchNorm2d(num_filters*4), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(num_filters*4, num_filters*8, kernel_size=4, stride=1, padding=1, padding_mode='reflect', bias=False), nn.BatchNorm2d(num_filters*8), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(num_filters*8, 1, kernel_size=4, stride=1, padding=1, padding_mode='reflect'), ) def forward(self, input_img, target_img): # 将输入和条件在通道维上拼接 x = torch.cat([input_img, target_img], dim=1) return self.model(x)

需要注意,判别器的输入是“成对拼接图”。把输入x和真实目标y拼接是正样本,输入x和生成结果G(x)拼接是负样本。这里用padding_mode='reflect'替代默认的零填充,能让边缘区域的感受野更自然,减少边缘伪影,这个技巧是实践中验证过有效的。

3.3 从零搭一个可训练的最小模型

网络结构写好了,现在把它们组装起来。下面的代码实现了模型初始化和基本配置。权重初始化这里大有讲究:pix2pix原文用的是形如N(0, 0.02)的高斯初始化,这种初始化能让网络在训练早期处于一个比较均衡的起始点。如果不做初始化或者用默认的PyTorch初始化,生成器和判别器的输出分布可能一开始就差异悬殊,导致训练不稳定。

def init_weights(net, init_gain=0.02): for m in net.modules(): if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d): nn.init.normal_(m.weight.data, 0.0, init_gain) elif isinstance(m, nn.BatchNorm2d): nn.init.normal_(m.weight.data, 1.0, init_gain) nn.init.constant_(m.bias.data, 0.0) # 实例化并初始化 generator = UnetGenerator(in_ch=3, out_ch=3, num_filters=64) discriminator = PatchDiscriminator(in_ch=3, num_filters=64) init_weights(generator) init_weights(discriminator)

跑一下看参数量,U-Net生成器大概在54M左右,PatchGAN判别器大概在2.7M左右,两者加起来60M左右的参数量。这个体量在GPU上训练是比较轻松的,一张中端显卡就能扛住。如果觉得生成器太大,可以把num_filters减到32,速度会明显提升,代价是生成的细节略有下降。我建议入门阶段先用默认的64,跑通之后再回头调。

4. 损失函数组合与训练循环

4.1 损失函数拆解:L1加GAN的化学反应

pix2pix总损失由两部分构成:GAN对抗损失加上L1像素损失。写成公式就是:

G损失 = GAN_loss(G) + λ * L1_loss(G)

GAN损失让生成图像“像真的”,L1损失让生成图像“靠近目标”。少了L1损失的模型虽然能骗过判别器,但生成结果可能跟输入完全不相关,因为它只追求真实性不管对齐。少了GAN损失的模型就是纯U-Net回归,前面说过会糊。两者配合的化学反应在于:GAN负责“像”,L1负责“准”。

权重λ为什么取100?这是个经验值,原文实验发现100的时候在“真实感”和“准确度”之间平衡最好。λ太大,生成图像会过于接近目标平均值,纹理细节被压制;λ太小,生成图像可能非常有真实感,但跟输入对不上。我在实际调参过程中试过λ=10和λ=500,前者生成器完全放飞,语义对应错乱;后者生成结果锐利度下降明显,线条模糊。没有绝对最优的λ,跟数据集有关系,但100作为起点基本不会出大问题。

另一个关键细节:L1损失为什么优于L2?L2损失对离群像素的惩罚是平方级别的,一张真实图像上偶尔出现的高亮噪点会把梯度拉向奇怪的方向,模型会变得保守,更倾向于输出模糊的平均值来减小惩罚。L1损失的梯度恒定,对边缘和细节的保留明显更好。这跟图像去噪领域的情况有点像,L1对结构保持更友好。

下面是训练循环里真正的损失计算代码。注意生成器要先更新,判别器后更新,顺序不能搞反。原因后面细说。

# 初始化损失函数 gan_loss_fn = nn.BCEWithLogitsLoss() l1_loss_fn = nn.L1Loss() lambda_l1 = 100 # 优化器 g_optimizer = torch.optim.Adam(generator.parameters(), lr=2e-4, betas=(0.5, 0.999)) d_optimizer = torch.optim.Adam(discriminator.parameters(), lr=2e-4, betas=(0.5, 0.999))

4.2 优化器与学习率策略

优化器选了Adam,但注意betas参数用的是(0.5, 0.999)而不是默认的(0.9, 0.999)。原因是生成对抗训练里的损失面非常不平滑,默认的一阶矩衰减系数0.9会把历史梯度的惯性带得太大,导致参数更新震荡。0.5的衰减让优化器更关注近期梯度,训练过程更稳定。这个改动是GAN训练里的公开经验,值得直接沿用。

学习率取2e-4,是GAN训练中很常用的初始值。如果训练到后期发现损失不再下降,可以按epoch数做线性衰减。比如总共200轮,前100轮固定学习率,后100轮从2e-4线性降到0。这样训练前期可以快速探索参数空间,后期逐渐收敛到更精确的点。实现的时候可以用torch.optim.lr_scheduler.LambdaLR,也可以手写一个step函数每轮更新一下优化器状态。

我自己的经验是,如果数据集很小(几千张图),线性衰减效果会更明显,因为小数据集上模型本来就容易过拟合,后期降低学习率能稳定生成质量。

4.3 训练循环的核心步骤

写训练循环之前,脑子里先过一遍整体流程,避免代码写得混乱。一次迭代要做的事情是:

  1. 取一对输入和真实目标图像
  2. 生成器算出假图像
  3. 用假图像骗判别器,计算生成器对抗损失和L1损失
  4. 更新生成器参数
  5. 用一次冻结的假图像,让判别器同时看真假两类样本,计算判别器损失
  6. 更新判别器参数

为什么先更新G再更新D?因为G更新后产生的假图像跟D之前看到的假图像分布不同,如果先更新D再更新G,D拿到的是G更新前的假图像,梯度会有滞后。先更新G、再更新D,能让D看到最新的假图像分布,判别器学得更准确。

下面是完整的单迭代训练代码。这里有一个细节:判别器计算正样本损失和负样本损失时,如果显存不够,可以分两段前向计算,先算fake的对半,再算real的对半,梯度积累后再优化器更新。

def train_one_step(input_img, target_img): # 清空梯度 g_optimizer.zero_grad() d_optimizer.zero_grad() # 1. 生成假图像 fake_img = generator(input_img) # 2. 生成器损失:尽可能骗过判别器 + 像素对齐 g_pred_fake = discriminator(input_img, fake_img) g_gan_loss = gan_loss_fn(g_pred_fake, torch.ones_like(g_pred_fake)) g_l1_loss = l1_loss_fn(fake_img, target_img) * lambda_l1 g_loss = g_gan_loss + g_l1_loss g_loss.backward() g_optimizer.step() # 3. 判别器损失:真实目标尽量判真,生成图像尽量判假 d_optimizer.zero_grad() fake_img = generator(input_img).detach() d_pred_real = discriminator(input_img, target_img) d_loss_real = gan_loss_fn(d_pred_real, torch.ones_like(d_pred_real)) d_pred_fake = discriminator(input_img, fake_img) d_loss_fake = gan_loss_fn(d_pred_fake, torch.zeros_like(d_pred_fake)) d_loss = (d_loss_real + d_loss_fake) * 0.5 d_loss.backward() d_optimizer.step() return g_loss.item(), d_loss.item(), fake_img

第二行的zero_grad只清了一次,但其实生成器更新完以后应该再清一次判别器的梯度。上面的代码里我在生成器step之后、判别器前向之前又单独调了一次d_optimizer.zero_grad(),这样逻辑更清晰。如果你偷懒在开头只清一次,会让判别器把生成器上一轮的梯度混进去,更新方向被污染,虽然不至于完全不能训练,但会减慢收敛。

注意第18行:fake_img = generator(input_img).detach()。detach非常重要。如果不detach,梯度会从判别器流向生成器,造成两个网络在反向传播时纠缠在一起,生成器会被重复更新两次,训练就失衡了。判断的标准其实很简单:更新判别器时,生成器的输出应当被视为不可训练的常量。

4.4 可视化与模型保存

训练GAN最忌讳闷头训练不看效果。每训练若干步,就应该拿测试集的输入喂给生成器,把输入、生成结果、真实目标三张图并排展示。这一步能让你很直观地看到模型有没有在学、学到什么程度。

我的个人建议是每100步打印一次损失,每500步做一次可视化。可视化时不要只放单张图,放一个2×3的网格对比不同样本,这样能暴露模型在部分数据上的过拟合或模式坍塌问题。

import torchvision import matplotlib.pyplot as plt def visualize(input_img, fake_img, target_img, save_path='sample.png'): images = torch.cat([input_img.cpu(), fake_img.cpu(), target_img.cpu()], dim=-1) grid = torchvision.utils.make_grid(images, nrow=1, normalize=True, range=(-1, 1)) plt.figure(figsize=(12, 4)) plt.imshow(grid.permute(1, 2, 0).numpy()) plt.axis('off') plt.savefig(save_path, bbox_inches='tight') plt.close()

模型checkpoint保存也有讲究。建议每5个epoch保存一次,保存内容不仅包括生成器和判别器的权重,还要保存优化器状态和当前epoch数。万一机器中断,可以从checkpoint恢复继续训练。我在项目里通常这样存:

checkpoint = { 'epoch': epoch, 'generator_state': generator.state_dict(), 'discriminator_state': discriminator.state_dict(), 'g_optimizer_state': g_optimizer.state_dict(), 'd_optimizer_state': d_optimizer.state_dict(), } torch.save(checkpoint, f'checkpoint_epoch_{epoch}.pth')

加载的时候逆序恢复就行了。有人只保存生成器权重,测试没问题,但如果想继续训练就麻烦,还得重新从头跑。

5. 常见问题与排查技巧实录

5.1 训练不收敛或损失震荡

GAN训练最常见的噩梦就是损失曲线像心电图一样狂跳。pix2pix因为有L1损失托底,通常比普通GAN稳定,但也不排除崩的可能性。

一个典型迹象是判别器损失快速降到接近0,生成器损失居高不下。这说明判别器太强,生成器完全没有反抗能力。解决思路有两个方向:一是削弱判别器,比如降低判别器学习率到生成器的三分之一;二是增强生成器,比如把生成器里的dropout关掉。实战中我遇到最多的原因是学习率设置过高,降到1e-4之后情况明显好转。

反过来,如果生成器损失一直降低但生成图像质量差,说明生成器在利用L1损失“作弊”,只输出目标图的模糊版本骗过L1,却没在对抗损失上学到真实纹理。这时候适当调低λ,让对抗损失占据更大权重,效果立竿见影。

5.2 生成结果模糊或过度平滑

生成图看起来“雾蒙蒙”的,先检查L1权重λ是不是太大。λ=100时偶尔也出现轻度模糊,尤其在那些输入本身纹理丰富的任务上。可以把λ降到50或40试试,代价是结果可能更“锐”但稳定性下降。

另一个常见原因是判别器感受野太小,只关注局部纹理,忽略了全局结构。如果你用的输入尺寸偏大(比如512×512),默认的70×70感受野不够看全局,可以考虑加一层卷积让感受野增大到140×140左右。但要小心,感受野过大会让判别器更关注全局,局部细节生成又有可能变弱,需要自己权衡。

还有一个偏门但经常被忽略的原因:生成器输出层用Tanh,如果数据归一化时用了ImageNet的mean/std而不是[-1,1],输出和目标的数值范围对不上,模型会一直挣扎但生成结果看起来“发灰发白”。检查一下数据归一化那一段,确认real image和fake image处于同一范围。

5.3 显存不足与训练速度优化

8GB显存跑256×256输入、batch size=1完全没问题。真正的杀手是可视化时同时forward多张测试图,GPU峰值会突然蹿高。建议可视化时把图搬到CPU再forward,或者一次只跑一张。

如果显存真的不够,优先缩小输入尺寸,比如从256降到192,而不是缩小batch size。因为pix2pix每个batch只有1-2张图,batch size再降已经没什么空间了。减小输入尺寸带来的显存节省是平方级的,效果明显。

另外一个容易被忽略的显存杀手是torch.cat拼接。判别器输入是输入图和目标图在通道维拼接,如果输入是3通道,拼接后就是6通道,中间层的特征图显存占用自然翻倍。这条不是让你改实现,而是提醒你别惊讶为什么判别器显存比想象中高。

5.4 模式坍塌与颜色偏移问题

模式坍塌在pix2pix里不如普通GAN那么严重,因为L1损失保证输出跟输入相关性很强,但确实会出现“输出千篇一律”“颜色偏灰”这类问题。

颜色偏灰的典型原因我前面提过:L1损失鼓励像素平均,导致模型倾向于输出安全的中间色。除了调低λ,还可以考虑把L1换成感知损失或者增加一个颜色直方图匹配约束。不过这些都是锦上添花的操作,先跑通原生pix2pix再说。

模式坍塌的一个隐蔽原因是batch size=1时BatchNorm失效。因为每个batch只有一张图,BatchNorm统计量是当前单张图的均值方差,抖动极大。如果发现训练曲线看起来正常但生成效果不稳定,可以试试用InstanceNorm替换所有BatchNorm。InstanceNorm在图像生成任务里很常用,对风格迁移特别友好,pix2pix换掉之后往往能稳定不少。代价是训练时间稍微变长,但表现力提升明显。

5.5 关于训练时长和硬件的一点个人体会

入门级的pix2pix训练,在单张中端显卡上,256×256分辨率、200个epoch、几千张训练图,大约需要6到10个小时。前50个epoch生成结果会从纯噪音变成有点形状,100个epoch左右才能看到一些像样的纹理细节,最后的50个epoch主要是在精修。别急着在50个epoch停止,我见过太多人模型还没训够就开始怀疑结构出问题。拿生成器输出和真实目标丢到一块板上对比,如果高频纹理已经明显接近,就可以停了。

如果时间特别紧,可以先用128×128分辨率做一次快速验证,确认代码和数据管线没问题,再换成256×256正式开训。这一步能省掉大量试错时间。

最后再分享一个我自己的习惯:训练结束后,不只看生成的单张图,会把整个测试集的输出来一遍,挑出效果最差的那几张看看它们跟输入的关系。通常你会发现自己模型的短板集中在某类特殊输入上,比如强光照、大阴影、复杂遮挡。找到规律后再针对性地增强该类数据,比盲目堆epoch数效率高得多。这就是我在pix2pix这个项目上摸爬滚打得到的最大教训:模型结构只是骨架,数据和调参才是让骨架活起来的血液。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询