资讯详情 U-Net与Pix2Pix图像转换实战:PyTorch从零实现与调参经验
📅 2026/10/21 4:03:45
像素级的图像转换image-to-image translation是深度学习里最直观、也最容易出效果的一类任务给你一张线条草稿生成上色后的成品给你一张语义分割图生成真实场景给你一张灰度照片自动补全颜色。做这些任务最经典的两个模型就是U-Net和Pix2Pix。U-Net最初是医学图像分割的标配骨架Pix2Pix则把它拿来当生成器配上PatchGAN判别器搭成了一个完整的条件GAN框架。这篇文章我会把两个模型的原理拆开讲透给出用PyTorch从零实现的完整代码最后分享我在实际训练中总结的调参经验和踩过的坑。适合已经能跑通简单图像分类或GAN入门项目、但想真正理解生成模型内部结构的读者。看完之后你不仅能自己复现Pix2Pix还能明白网络结构为什么要这么设计。1. 为什么我还在坚持从零实现这两个模型1.1 框架能帮你跑通但帮不了你变通我刚接触图像生成任务的时候第一反应也是找现成仓库。网上确实有大量封装好的U-Net和Pix2Pix实现安装依赖、准备数据、微调几步就能在自己的数据上跑出还看得过去的成果。但很快我发现一个问题当我想改一点功能比如把输出从单通道改成多通道或者把判别器从PatchGAN换成全局判别器我根本不知道该从哪下手。现成代码对我来说是个黑盒能跑但改不动。从零实现的价值首先在于它逼你把每一张张量的维度变化走一遍。你学再多的结构图不如亲手写一次forward函数看着一个3x256x256的输入经过卷积、池化、跳跃连接、上采样最终变成3x256x256的输出。这个过程中你才会真正明白通道数什么时候翻倍、什么时候减半跳跃连接拼接发生在哪个维度上采样之后特征图尺寸怎么恢复。这些知识是任何现成代码都代替不了的。说白了框架能帮你跑通流程但帮不了你在关键时刻做变通。1.2 从零实现练的是故障诊断能力训练生成模型跟训练分类模型是两种完全不同的体验。分类模型的loss下降基本意味着精度提升方向很明确。生成模型则不然loss降了不代表生成质量好甚至经常出现判别器loss极端而生成器输出一片糊的情况。遇上这种问题你只能靠对架构的深刻理解去定位——是判别器太强了L1损失权重低了还是BatchNorm在生成器里偷偷引入了批次依赖我自己手写过一遍U-Net和Pix2Pix之后再读各种变体论文明显轻松得多。因为我知道哪些部分是通用骨架哪些是某篇论文特有的巧思。这篇文章接下来要讲的就是我实现时的完整思路、代码和排错记录你可以照着写一遍然后按自己的理解去改。2. U-Net结构逐层拆解编码器、解码器与跳跃连接2.1 收缩路径图像是怎么一步步被榨干的U-Net的输入是一张完整图像比如3通道、256x256。收缩路径的工作是用一层层卷积把空间分辨率降下来同时把通道数提上去。每一层都是两次卷积一次最大池化通道数按64、128、256、512这样翻倍特征图尺寸每次缩小为原来的一半。为什么这么设计一个很关键的直觉是在图像任务里小卷积核只能看到局部区域想捕捉更大范围的语义信息比如这是一辆车而不是一堆像素必须有足够大的感受野。下采样就是在干这件事分辨率减半之后同样大小的卷积核覆盖的实际区域变成了原来的四倍。通道数增加则是为了给每个位置存放更丰富的语义编码。到最深层的时候特征图可能只有16x16甚至8x8但每个像素代表的是输入图上的一大块抽象语义。用压缩文档来做类比收缩路径像是把一份高清资料反复压缩成越来越小的缩略图缩略图保留了整体结构但细节没了。问题在于后面的解码器怎么把细节找回来。2.2 扩展路径与跳跃连接细节是如何被救回来的解码器要做的事情刚好和编码器相反把低分辨率、高通道数的特征图逐步恢复成高分辨率、低通道数的输出。如果只是简单地把缩略图上采样回原尺寸丢失的空间细节是找不回来的。U-Net的解决办法就是跳跃连接把收缩路径每一层下采样之前的特征图直接拼接到解码器对应层的输入上。这里有个细节值得多说一句为什么是拼接而不是相加拼接是沿通道维度把两组特征图叠在一起解码器自己学习该从底层细节特征里取多少、从高层语义特征里取多少相加则相当于提前定死了融合比例两组信息互相干扰的可能性更大。原论文选的是拼接实践里证明效果也更好。还拿压缩文档类比把高清照片压缩成缩略图再放大细节永远回不来。但如果放大时手里还有原始照片的局部裁剪版做参考你就能对照着把边缘和纹理补回去。跳跃连接就是递给解码器的那份参考原件。这也是U-Net在分割、生成这类像素级密集预测任务上比普通编解码结构强的根本原因高层语义和低层细节被同时利用了。2.3 为什么是3x3卷积加2倍下采样U-Net全篇几乎只用3x3卷积输出用1x1卷积调节通道数下采样固定用2x2、stride为2的最大池化。这套配置已经成为图像分割和生成领域的默认习惯背后是有原因的。3x3卷积加padding1卷积前后特征图尺寸不变这让网络深度可以随意堆叠而不需要操心空间尺寸变化。两次3x3卷积叠加相当于一次5x5卷积的感受野但参数量更少还引入了非线性表达能力更强。2倍下采样则是精度和效率的折中下采样太多细节丢失严重太少感受野增长太慢深层语义提取不出来。4次2倍下采样对应16倍的压缩在256x256输入上正好把特征图压到16x16这个深度在实际任务中表现最为均衡。3. 用PyTorch从零搭一个U-Net3.1 基础模块双卷积块的设计我把U-Net拆成三个部分双卷积块、编码器列表、解码器列表。先写最基础的双卷积块它是整个网络唯一重复使用的积木。import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_channels, out_channels): super(DoubleConv, self).__init__() self.conv nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue), nn.Conv2d(out_channels, out_channels, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x)两个设计细节我踩过坑提醒一下。第一卷积的bias设置为False。因为后面紧跟BatchNormBatchNorm会做平移操作卷积层的偏置是冗余的留着反而可能造成数值不稳定。第二BatchNorm放在卷积和激活之间。这个顺序是主流做法——先归一化再激活能避免激活函数输出的分布偏移被放大。如果你用的是GNGroupNorm而不是BatchNorm位置同理。3.2 编码器、瓶颈与解码器的拼装有了双卷积块U-Net主体就是把积木按顺序搭起来。我把通道数序列设成元组方便后面改不同规模的网络。class UNet(nn.Module): def __init__(self, in_channels3, out_channels3, features(64, 128, 256, 512)): super(UNet, self).__init__() self.downs nn.ModuleList() self.ups nn.ModuleList() self.pool nn.MaxPool2d(kernel_size2, stride2) # 编码器 for f in features: self.downs.append(DoubleConv(in_channels, f)) in_channels f # 瓶颈层 self.bottleneck DoubleConv(features[-1], features[-1] * 2) # 解码器 for f in reversed(features): self.ups.append(nn.ConvTranspose2d(f * 2, f, kernel_size2, stride2)) self.ups.append(DoubleConv(f * 2, f)) self.final_conv nn.Conv2d(features[0], out_channels, kernel_size1) def forward(self, x): skip_connections [] for down in self.downs: x down(x) skip_connections.append(x) x self.pool(x) x self.bottleneck(x) skip_connections skip_connections[::-1] for idx in range(0, len(self.ups), 2): x self.ups[idx](x) skip skip_connections[idx // 2] if x.shape ! skip.shape: x nn.functional.interpolate( x, sizeskip.shape[2:], modebilinear, align_cornersFalse ) x torch.cat((x, skip), dim1) x self.ups[idx 1](x) return self.final_conv(x)encoder循环里每个DoubleConv的输出会被存进skip_connections然后才做池化。所以跳跃连接里存的是下采样之前的特征图这是U-Net的关键千万别搞反。瓶颈层我把通道数翻倍在[64,128,256,512]的配置下就是512进、1024出。decoder部分每个阶段先做转置卷积上采样通道数减少一半然后和对应的跳跃连接拼接通道数又回到一倍再做双卷积降回目标通道数。为什么要用nn.ModuleList而不是Python列表因为ModuleList里的模块会被注册进模型的参数表这样.to(device)、.state_dict()才能覆盖到所有层。直接用Python列表存模块的话参数根本不会被记录模型训练时那些层的权重纹丝不动。这个坑非常隐蔽我在初学的时候查了很久才找到原因。3.3 forward里的张量维度变化我习惯在第一次跑之前手动把每个阶段的张量形状在纸上写一遍。以3x256x256输入、features(64,128,256,512)为例阶段输入形状输出形状第一个双卷积3x256x25664x256x256池化64x256x25664x128x128第二个双卷积64x128x128128x128x128池化128x128x128128x64x64第三个双卷积128x64x64256x64x64池化256x64x64256x32x32第四个双卷积256x32x32512x32x32池化512x32x32512x16x16瓶颈512x16x161024x16x16解码第一层1024x16x16拼512512x32x32解码第二层512x32x32拼256256x64x64解码第三层256x64x64拼128128x128x128解码第四层128x128x128拼6464x256x256解码时转置卷积的kernel_size2、stride2刚好让空间尺寸翻倍所以每个阶段都能和对应跳跃连接对齐。你会发现解码每一层的输入通道数都是2倍一半来自上采样特征一半来自跳跃连接。这就是为什么转置卷积输出是f、拼接后变成2f再用DoubleConv把2f降回f。整个流程逻辑严密只要能把这个表格推出来U-Net你就吃透了大半。4. Pix2Pix的对抗框架不是所有GAN都适合做图到图4.1 条件生成判别器到底在看什么Pix2Pix本质上是一个条件GANcGAN。普通GAN的生成器输入一个随机噪声输出一张图生成什么内容完全不受控制。Pix2Pix做了一个关键改动生成器除了噪声还接收一张条件图——你要转换的源图。比如语义分割图转街景分割图就是条件图生成器的任务是在它的约束下生成街景图。配套地判别器的输入也变了。它不再只看生成的或真实的目标图而是看一个拼接对条件图生成图或者条件图真实目标图。判别器的任务是判断这个拼接对是假的还是真的。为什么要连条件图一起看因为只判断单张图的真假生成器很容易生成一张看起来真实但完全不管条件图的图。把条件图绑进去判别器被迫检查生成图和条件图之间的一致性生成器也就被迫在保留条件信息的前提下去骗过判别器。这个设计带来的效果很直观如果你做过无条件的GAN你会发现生成结果是全局上像目标分布但局部对应关系完全混乱。改成条件生成后输入输出在空间结构上严格对齐这正是图像转换任务想要的。4.2 PatchGAN为什么比全局判别器好用Pix2Pix的判别器不是把整张图压成一个0到1的标量而是输出一个特征图——比如16x16的矩阵矩阵里的每个值代表原图一个局部patch的真假评分。这就是PatchGAN原文用的patch_size是70x70左右的感受野。为什么局部判别器效果更好一个全局判别器要判断整张图真假语义太全局很容易被大结构特征比如整体色调、构图糊弄过去却对纹理细节不敏感。PatchGAN把注意力放在局部一致性上生成器必须在每个局部区域都做得足够逼真才能骗过它。同时每个patch独立打分相当于给生成器提供了更丰富的梯度信号训练更稳定还大幅减少了参数量。有一件事值得强调PatchGAN并不是把图像真的切成小块分别喂进去它就是在标准卷积神经网络上做卷积利用感受野天然覆盖局部区域最终经过若干层卷积后输出平滑的patch概率图。所以实现起来比想象中简单得多下一节直接上代码。4.3 L1损失让生成器别只顾着骗人如果只用对抗损失训练生成器一个典型毛病是生成结果漂移整体风格像目标域但和输入条件图对应的结构细节对不上比如不该突出的纹理被乱加了。Pix2Pix解决这个问题的办法很朴素给生成器加一个重建约束让它输出的图和真实目标图在像素级别上尽量接近。原论文用的是L1损失而不是L2并且权重给到了100。为什么选L1经验上L1损失对边缘更友好L2损失会把误差平方放大模型倾向于输出模棱两可的平均色结果就是图像变糊。L1对异常值的惩罚温和一些生成结果边缘更锐利。权重100意味着对抗损失只起锦上添花的作用L1才是主约束——它保证了生成图大体正确对抗损失负责让细节看起来像真的。我在实践中也验证过lambda取10时生成图明显更随意取1000时又有点过度平滑100确实是个好用的默认值。5. 从零写PatchGAN判别器与训练主循环5.1 PatchGAN判别器代码逐行说明PatchGAN的代码比生成器简单得多本质就是一个纯卷积网络输入是拼接后的6通道图3通道条件图3通道目标图或生成图输出是patch概率图。class PatchGAN(nn.Module): def __init__(self, in_channels6): super(PatchGAN, self).__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.BatchNorm2d(128), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv2d(128, 256, kernel_size4, stride2, padding1), nn.BatchNorm2d(256), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv2d(256, 512, kernel_size4, stride2, padding1), nn.BatchNorm2d(512), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv2d(512, 1, kernel_size4, padding1) ) def forward(self, x): return self.model(x)几个细节说一下。所有卷积的kernel_size4前四层stride2做下采样最后一层stride1保住空间分辨率。激活函数用LeakyReLU而不是ReLU负斜率0.2——判别器需要保留小梯度信号ReLU把负数区域直接清零容易让训练失去信息。第一个卷积层后面没有BatchNorm这是从DCGAN那里沿袭下来的实践第一层直接用原始像素输入过早归一化反而会破坏低层特征的分布实验下来效果更稳定。以256x256输入为例stride2的卷积连续做四次特征图变成16x16最后一层stride1再输出16x16的patch评分图。在小的输入尺寸上比如128x128输出会变成8x8patch更小对高频细节的要求相对降低这在特定任务上可能是优点。5.2 训练循环里容易被忽略的顺序问题写训练循环最需要注意的是更新判别器用的假图必须来自冻结状态的生成器。我刚开始实现时直接在训练判别器时用了生成器新产生的图结果反向传播把生成器的梯度通过假的路径也传了一遍导致两个网络的梯度纠缠在一起loss直接乱套。正确做法是fake generator(source).detach()截断梯度。def train_one_epoch(generator, discriminator, dataloader, opt_g, opt_d, l1_loss, bce_loss, lambda_l1100): for source, target in dataloader: source source.to(device) target target.to(device) # 第一步训练判别器 fake generator(source).detach() d_real discriminator(torch.cat((source, target), dim1)) d_fake discriminator(torch.cat((source, fake), dim1)) loss_d bce_loss(d_real, torch.ones_like(d_real)) \ bce_loss(d_fake, torch.zeros_like(d_fake)) opt_d.zero_grad() loss_d.backward() opt_d.step() # 第二步训练生成器 fake generator(source) d_fake discriminator(torch.cat((source, fake), dim1)) loss_g_adv bce_loss(d_fake, torch.ones_like(d_fake)) loss_g_l1 l1_loss(fake, target) * lambda_l1 loss_g loss_g_adv loss_g_l1 opt_g.zero_grad() loss_g.backward() opt_g.step()生成器第二步为什么不detach因为对抗损失需要梯度流经判别器再回到生成器这时候生成器是要被更新的所以不能截断。判别器损失里d_fake那一项如果忘了把fake detach虽然短期看不出来但训练几十轮后会发现生成器更新方向和预期不一致定位特别麻烦。我建议写训练循环时把两步更新的边界划清楚第一步只更新判别器第二步只更新生成器。先在纸上写明每一步哪些网络需要梯度再动键盘。这个习惯能替你省下大量调试时间。6. 训练配置与参数调优的实测经验6.1 数据准备成对数据要怎么喂Pix2Pix要求成对数据——同一张场景的源图和目标图。常见的做法是把两张图左右拼在一起存成一张图训练时按宽度切半。我建议在数据读取阶段直接返回两个tensor不要真的拼图再切因为翻转和裁切等数据增强在两个tensor上要保证同步。数据增强方面原论文的做法是先把图片缩放到286x286再随机裁切到256x256最后做随机水平翻转。这个先放大再随机裁剪的组合非常关键它让模型对平移和缩放有一定的鲁棒性比直接固定尺寸训练效果好不少。注意生成数据和目标数据要用同一组随机种子做同样的变换否则对应关系就错位了。像素归一化必须统一到[-1,1]也就是(x / 127.5) - 1。这个细节很多人忽略。生成器最后的激活函数一般是Tanh输出天然在[-1,1]判别器看到的真实图如果不是同一个范围两个网络对数值尺度的理解就不一致训练会异常困难。6.2 学习率、Batch Size与轮数的取舍我训练Pix2Pix的默认配置如下超参数常用取值说明优化器Adam两个网络都用学习率2e-4原论文默认稳妥betas(0.5, 0.999)第一项取0.5稳定性更好lambda_L1100L1约束权重batch size1~8256x256下受显存限制训练轮数100以上生成质量随轮数缓慢提升Batch Size这块我要多说一句。CV里的分类任务batch size越大越好但GAN训练不是这个逻辑。256x256分辨率的图一张就要占不少显存而且batch size太大会让判别器每步见过太多样本梯度估计过于平稳反而容易把生成器逼进只输出以偏概全的假图的局部陷阱。我实际用的最多的就是batch size1也就是所谓online learning配合BatchNorm反而有一种独特的正则化效果。BatchNorm在batch size1时退化成InstanceNorm但实验下来依然能收敛可见Pix2Pix对BatchNorm的依赖没有想象中那么强。训练轮数上别指望像分类模型那样三十轮就出结果。图像生成任务通常在80轮之前生成质量都在缓慢爬升后面每20轮还能看到可感知的细节提升。没有特殊原因我建议至少训练100轮再下结论。6.3 两个网络学习节奏不同步怎么办GAN训练里最常见的失衡现象是判别器学得太快D loss很快降到接近0生成器再怎么练都骗不过它。原因是判别器任务本来就比生成器简单——判断真假是一个二分类而生成一个逼真图像是困难的生成问题。遇到D loss极端偏低我有几个实测有效的处理手段。第一降低判别器的学习率比如把D的lr设置为G的1/5不推荐修改betas。第二把真实标签从1改成0.9或随机0.9到1.0之间也就是label smoothing降低判别器的绝对自信度。第三给生成器的对抗损失换个非饱和形式也就是让生成器优化最大化D对假图的预测错误率而不是最小化D对假图的正确率对应到代码就是G的loss目标是torch.ones_like(d_fake)而不是torch.zeros_like(d_fake)我上面的代码用的就是这种。如果发现反过来G太强、D跟不上生成图会开始出现色调异常那就反过来提高D的lr或者把lambda适当调小。7. 从零实现最容易踩的五个坑7.1 通道数没对齐运行时才报错U-Net里最容易写错的是解码器的通道数。很多人记住了跳跃连接拼接却忘了拼接之后通道数翻倍导致DoubleConv的输入通道声明和实际tensor对不上。这属于框架能检测出来的错误报错信息会提示shape mismatch好在不算难定位。更隐蔽的是某些层之间通道数恰好碰巧一致只是语义错了——比如把应该拼接的skip连接加到了错误层级输出看起来正常但模型性能始终上不去。排查方法很简单每层forward后打印一次张量形状和我在3.3节给的表逐行对照。7.2 输入尺寸必须被下采样次数整除如果输入图是300x300经过四层2倍下采样会得到18.75PyTorch会直接给你报错。所以输入尺寸必须是16的倍数4次下采样或32的倍数5次下采样。有人会用resize强行改尺寸但要注意把不规则数据硬拉成固定尺寸可能有信息变形。更优雅的做法是在数据加载阶段做居中裁剪或padding然后在计算loss时只对有效区域做L1损失。7.3 棋盘格伪影棋盘格伪影是转置卷积的经典问题尤其在低分辨率特征上采样时明显。转置卷积是插值卷积的组合kernel_size2、stride2时卷积核的感知区域在边界处重叠不均匀就会产生格子状纹路。这是我在生成灰度图时最常看到的瑕疵。解决办法有两个。一个是改用nn.Upsample(modebilinear)加普通卷积的组合替代转置卷积生成器结构改成上采样到目标尺寸再做3x3卷积效果接近且没有棋盘格。另一个是用kernel_size4、stride2、padding1的转置卷积感受野分布更均匀也能缓解问题。我倾向于第二种因为改动最小且保留转置卷积的可学习性。7.4 判别器坡度太陡生成图灰蒙蒙如果生成器的输出色调整体偏灰、对比度低、边缘模糊多半是L1权重过大把生成器压得太死。生成器发现只要输出接近目标像素的平均值L1损失就能降得很低于是它选择猥琐发育对抗损失带来的锐化信号根本不足以让它冒险。我在lambda从100改成200时遇到过这个问题回退到100并配合label smoothing之后就恢复了。另外一个小技巧观察生成图如果整体偏灰还可以检查生成器输出端的激活函数。如果误用了Sigmoid输出会被限制在[0,1]且后期梯度饱和换成Tanh配合[-1,1]归一化会好很多。7.5 过拟合还是欠训练先分清再动手小数据集上很容易看到训练几百步后生成图像像模像样但验证集上效果崩掉——这是典型的过拟合。生成模型同样过拟合而且更隐蔽GAN没有类似分类准确率的指标你很难通过数值判断。我的建议是训练过程中定期把固定的一批源图喂给生成器保存输出成网格图人工盯生成趋势。输出稳定且不崩说明训练健康输出在某个阶段突然变得千篇一律八成是模式崩溃或判别器失衡。这套人工评估流程不用写代码就是训练循环里每N轮save一次图像事后按时间顺序翻看。它能告诉你的信息比一堆loss曲线多得多。最后再分享一条个人经验从零实现U-Net和Pix2Pix的过程里真正让你成长的未必是跑通的那一刻而是那个输出全黑、loss全是NaN、维度怎么都对不上的深夜。你会被迫去逐层打印张量、去重读论文里的每个设计细节、去理解BatchNorm和归一化到底在做什么。把这些经历记录下来往后你再看任何生成模型的实现都会有一种这里为什么这么做的直觉。如果你正准备动手复现我的建议是别急着加各种花哨的模块先把最朴素的U-Net加PatchGAN跑通观察它的行为再慢慢往里面加注意力、加多尺度每一步都用自己的数据验证效果。这条路走完你对整个生成式视觉的理解会上一个台阶。