简介这套手写Unet工程面向深度学习新手以图像到图像生成为典型任务兼顾理论讲解与动手实践便于直接理解Unet模型结构、跳跃连接与像素级训练流程。RAR压缩包共614个文件、约47.71MB其中600张png多为输入输出样例与训练生成结果4个py覆盖数据预处理、模型搭建、训练及验证核心逻辑另有pth权重与iml、xml工程配置解压后即可按目录结构运行。目前已有792人学习/下载适合作为入门阶段的实践参照。通过这套代码可以完整走通图像读取、数据增强、模型训练到测试评估的全流程结合可视化输出图能直观对比生成效果还可基于已有权重修改网络层数、损失函数或超参数观察不同配置下的训练表现对快速上手Unet和图像翻译很有帮助。1. 为什么图像到图像生成偏偏绕不开Unet大多数教程把 U-Net 绑定在“医学图像分割”这一场景里这让不少新人误以为它只是一个分割专用工具。实际上图像去噪、图像上色、线稿转成品、风格迁移这类被称为“图像到图像生成”的任务里主流模型中的生成器骨架十有八九还是 U-Net。原因并不玄学图像到图像生成要求输入和输出在空间上严格对齐而 U-Net 的编码器-解码器结构天然适合这种“像素对像素”的映射。它不需要引入判别器也不需要复杂的对抗训练用最基本的 L1 损失就能把一张模糊图变清晰把灰度图染上颜色。对于小白来说自己手写一个 U-Net 并跑通一次图像到图像生成比直接调现成库更能建立起对卷积、特征图、上采样这些概念的实感。后面我按照实际调试这条流程的顺序来写从结构拆解、代码实现再到参数调整和验证。2. Unet结构拆解先认清那张U形图再动手写代码2.1 编码器、解码器与跳跃连接为什么它天生适合图像到图像生成U-Net 的名字来自网络结构画出来像字母 U。左侧是编码器由多次卷积和池化组成每一步都会让特征图变小、通道数变多右侧是解码器用转置卷积逐步恢复图像尺寸和通道数。没有跳跃连接时这个结构就是一个普通的自编码器输入被压缩到底层再还原出来。自编码器的问题在于经过多次池化后高频细节被大量丢弃还原出来的图像很糊而且对边缘位置非常敏感。跳跃连接直接解决了这个问题。编码器每一层池化之前都会保存一份特征图解码器在对应层上采样之后把这份特征图沿通道维度拼接起来。这样一来解码器既能拿到底层的语义信息又能直接参考同尺寸的浅层细节。图像到图像生成里输入和输出是空间对齐的跳跃连接把输入细节旁路给输出端生成过程没必要从零猜测边缘和纹理。这也是 U-Net 比普通自编码器更适合做生成任务的根本原因。2.2 下采样与上采样的形状变化先搞懂张量怎么流在给代码之前先看清一张 256×256 的图走过网络时的尺寸变化。这里用 4 层下采样和 4 层上采样统一用 3×3 卷积、padding1、步长为 1所以卷积本身不改变宽高只有 MaxPool 和 ConvTranspose 才会改变宽高。阶段输入尺寸操作输出尺寸编码器第 1 层256×2563×3 conv ×2 pool128×128编码器第 2 层128×1283×3 conv ×2 pool64×64编码器第 3 层64×643×3 conv ×2 pool32×32编码器第 4 层32×323×3 conv ×2 pool16×16bottleneck16×16conv ×216×16解码器第 1 层16×16转置卷积 拼接32×32解码器第 2 层32×32转置卷积 拼接64×64解码器第 3 层64×64转置卷积 拼接128×128解码器第 4 层128×128转置卷积 拼接256×256需要注意输入尺寸必须是 2 的整数倍否则最深层会出现尺寸除不尽的情况。我一般统一缩放成 256×256省去 F.interpolate 兜底分支尺寸不匹配的问题几乎不会出现。2.3 最小可用Unet的PyTorch实现我手写时习惯把“两次卷积 归一化 激活”封装成一个 DoubleConv 块这个基础单元几乎所有 U-Net 变体都在用。import torch import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, kernel_size3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, kernel_size3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.conv(x)nn.Conv2d的padding1在卷积核为 3 时能保持特征图宽高不变这样网络只需要关心池化和上采样对尺寸的影响。BatchNorm 在通道维做归一化让每层输入的分布更稳定也能让训练收敛快不少。然后是 U-Net 主体class UNet(nn.Module): def __init__(self, in_channels3, out_channels3, features(64, 128, 256, 512)): super().__init__() self.pool nn.MaxPool2d(2) self.encoder nn.ModuleList() in_ch in_channels for f in features: self.encoder.append(DoubleConv(in_ch, f)) in_ch f self.bottleneck DoubleConv(features[-1], features[-1] * 2) self.up_convs nn.ModuleList() self.decoder nn.ModuleList() for f in reversed(features): self.up_convs.append(nn.ConvTranspose2d(f * 2, f, kernel_size2, stride2)) self.decoder.append(DoubleConv(f * 2, f)) self.out_conv nn.Conv2d(features[0], out_channels, kernel_size1) def forward(self, x): skips [] for enc in self.encoder: x enc(x) skips.append(x) x self.pool(x) x self.bottleneck(x) skips skips[::-1] for i, (upconv, dec) in enumerate(zip(self.up_convs, self.decoder)): x upconv(x) skip skips[i] if x.shape[2:] ! skip.shape[2:]: x F.interpolate(x, sizeskip.shape[2:], modebilinear, align_cornersFalse) x torch.cat([x, skip], dim1) x dec(x) return self.out_conv(x)features控制每个编码器层输出的通道数。默认值(64, 128, 256, 512)是 4 层下采样参数总量约 3000 万如果显存只有 6G可以改成(32, 64, 128, 256)对入门演示完全够用。nn.ConvTranspose2d的stride2会让输出宽高翻倍第 2.2 节表格里的尺寸变化正是由它带来的。torch.cat在通道维拼接上采样结果和跳跃连接的特征图这是跳跃连接真正起作用的代码位置。最后 1×1 卷积把解码器最后一层的输出通道数压缩成out_channels。这里没有加激活函数因为训练时损失函数直接作用在线性输出上测试时再按任务加 sigmoid 或 tanh比在模型内部固定激活更灵活。2.4 输入输出通道怎么定灰度图、RGB、多标签任务类型in_channelsout_channels典型案例灰度图到灰度图11去噪、超分灰度图到RGB13线稿上色RGB到RGB33风格转换、图像修复RGB到多标签图3类别数图像分割对于小白入门我建议先用 RGB 到 RGB比如“模糊原图到清晰原图”。这样输入输出都是三通道不需要额外处理掩码损失函数直接用 L1 就能跑通。做灰度图任务时注意把Image.open(...).convert(L)后的单通道图和网络定义的in_channels1对应上。out_channels 最常见的错误是把“生成彩色图”写成 1结果训练时目标图是三通道形状直接对不上报错也会指向维度不匹配。3. 用自己手写的Unet跑通第一轮图像到图像生成3.1 数据准备成对图像的Dataset 该怎么组织图像到图像生成和图像分类最大的不同在于每个训练样本必须是一对图输入图 x 和目标图 y。常见做法有两种。一种是把 x 和 y 横向拼接在同一张 PNG 里Dataset 里按宽度切成两半另一种是建立两个目录同名文件分别放入input/和target/。我更喜欢第二种文件管理清晰不会出现拼接边界带来的处理问题。第一步是把训练数据整理成下面的形式data/ input/ 001.png 002.png target/ 001.png 002.png两个目录下文件名必须严格一一对应。顺序不一致时用sorted(glob(...))排序只要两端文件名排序规则一致即可。import glob 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, image_size256): self.input_paths sorted(glob.glob(input_dir /*.png)) self.target_paths sorted(glob.glob(target_dir /*.png)) assert len(self.input_paths) len(self.target_paths), 输入和目标数量不一致 self.transform T.Compose([ T.Resize((image_size, image_size)), T.ToTensor(), T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]), ]) def __len__(self): return len(self.input_paths) def __getitem__(self, idx): x Image.open(self.input_paths[idx]).convert(RGB) y Image.open(self.target_paths[idx]).convert(RGB) return self.transform(x), self.transform(y)T.Normalize把三通道图像从 [0,1] 映射到 [-1,1]输入分布接近零均值梯度更平稳。我建议 image_size 至少取 256太小会让边缘细节丢失训练出来整张图都是糊的。如果要做数据增强必须把相同的随机变换同时应用到 x 和 y 上import random def __getitem__(self, idx): x self.transform(Image.open(self.input_paths[idx]).convert(RGB)) y self.transform(Image.open(self.target_paths[idx]).convert(RGB)) if random.random() 0.5: x T.functional.hflip(x) y T.functional.hflip(y) return x, y不要只翻转输入不翻转目标这样会破坏像素对应关系。注意不要把验证图片的文件名混在训练目录里否则 DataLoader 会把它当成训练数据导致验证结果虚高。3.2 训练循环的骨架损失、优化器、验证选损失函数时第一轮推荐用nn.L1Loss()。L2 损失会让 loss 更快地降到较小值但输出会变得模糊L1 训练出来的图片边缘更清晰更符合人的视觉判断。from torch.utils.data import DataLoader import torch.optim as optim device torch.device(cuda if torch.cuda.is_available() else cpu) model UNet(in_channels3, out_channels3).to(device) criterion nn.L1Loss() optimizer optim.Adam(model.parameters(), lr1e-4, betas(0.5, 0.999)) scheduler optim.lr_scheduler.ReduceLROnPlateau(optimizer, modemin, factor0.5, patience10) dataset PairedDataset(data/input, data/target) train_loader DataLoader(dataset, batch_size8, shuffleTrue, num_workers4, pin_memoryTrue) for epoch in range(100): model.train() total_loss 0.0 for x, y in train_loader: x, y x.to(device), y.to(device) pred model(x) loss criterion(pred, y) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() avg_loss total_loss / len(train_loader) scheduler.step(avg_loss) print(fepoch {epoch1}, loss {avg_loss:.4f})这段代码是标准的三件套前向计算 loss、反向传播、更新参数。Adam 的betas(0.5, 0.999)是生成任务里常用的配置把一阶矩衰减从默认 0.9 改成 0.5会让优化器更快地“忘掉”历史梯度防止开始阶段震荡太剧烈。如果 loss 曲线一直在跳把 lr 降到 3e-5如果 loss 下降太慢检查是不是没有开 CUDACPU 训练 256×256 的 U-Net 会慢到让人怀疑人生。每 5 个 epoch 保存一次模型权重torch.save(model.state_dict(), funet_epoch_{epoch1}.pth)3.3 跑一次看效果从logits到可视化图像的完整流程训练中只看 loss 数值是不够的必须把预测结果还原成图片看一眼。还原过程要把归一化反转即执行x * 0.5 0.5把 [-1,1] 拉回 [0,1] 后再交给torchvision.utils保存。import torchvision.utils as vutils def denorm(tensor): return tensor * 0.5 0.5 model.eval() with torch.no_grad(): sample_x, sample_y next(iter(train_loader)) sample_x, sample_y sample_x.to(device), sample_y.to(device) pred model(sample_x) vutils.save_image(denorm(sample_x), result_input.png, nrow4) vutils.save_image(denorm(sample_y), result_target.png, nrow4) vutils.save_image(denorm(pred), result_pred.png, nrow4)这里nrow4表示一行放 4 张图。如果 batch size 是 8输出就是两行。保存后打开result_pred.png第一轮大概率是一片偏灰的噪声这很正常训练到第 10 轮左右应该能看到与 target 接近的整体色调但边缘有重影第 30 轮之后细节会明显清晰起来。如果到 30 轮还是一片噪点通常是输入和目标没有对齐或者数据顺序错了。4. 把Unet从“能跑”调到“好用”参数、损失函数和典型坑4.1 损失函数怎么选L1、L2、感知损失图像到图像生成不是只要一个交叉熵就万事大吉。像素级生成任务里损失函数直接决定最终输出是锐利还是模糊。L2 损失对误差的惩罚是平方级模型会把错误尽可能分摊到各个像素上结果就是输出向所有可能结果的平均靠拢整体发灰。L1 损失对异常值更宽容能保留更多边缘信息。如果上一章用的是 L1继续往下调参时可以保留它。如果觉得 L1 生成的结果纹理不够自然再往上一步是加感知损失。感知损失的核心是让生成图与目标图在预训练 VGG 网络的特征空间里更接近而不是只在像素值上接近。我一般这样实现import torchvision.models as models class PerceptualLoss(nn.Module): def __init__(self): super().__init__() vgg models.vgg16(pretrainedTrue).features[:16].eval() self.vgg vgg.requires_grad_(False) self.l1 nn.L1Loss() def forward(self, pred, target): return self.l1(self.vgg(pred), self.vgg(target))组合使用时total_loss l1_loss(pred, target) 0.1 * perceptual_loss(pred, target)。要先把输入尺寸缩放到至少 224×224因为 VGG 的官方预训练权重在 ImageNet 的 224×224 输入上训练。感知损失权重不要一开始就设 0.5否则模型会过度关注高层语义而忽略颜色输出会出现类似卡通化的质感。4.2 学习率、归一化、数据增强对生成质量的影响训练曲线的震荡九成和学习率有关。U-Net 不是特别深的网络但同样对学习率敏感。Adam 默认 lr1e-3 对分类任务很好用放到像素级生成里经常会崩。我测试下来1e-4 是 256×256 图像上一个比较稳的起点当 batch size 从 8 降到 2 时建议继续下调到 5e-5因为小 batch 带来更大的梯度方差需要更小步长来稳定。关于归一化BatchNorm 在 batch size 大于 8 时表现稳定。如果因为显存只能跑 batch2那么每个 batch 的均值和方差跳动太大训练时看着收敛测试时却各种偏色。最简单的替代方案是把模型里的 BatchNorm 换成 InstanceNormnn.InstanceNorm2d(out_ch, affineTrue)InstanceNorm 对单张样本独立计算均值和方差不受 batch 大小影响在图像生成任务里是更常用的选择。对于小白入门保留 BatchNorm 也没问题但要知道这个陷阱。数据增强不要滥用。随机水平翻转、随机垂直翻转、随机裁剪都能提升泛化随机颜色抖动、随机灰度化会破坏输入和目标之间的语义一致性图像到图像生成里最好避免。裁剪尺寸最好控制在原图尺寸的 80% 到 100% 之间过小会切掉关键结构。4.3 Unet使用时的注意事项与常见问题排查表问题可能原因解决方式loss 正常下降输出全灰归一化被重复执行或输出层没有激活检查数据 pipeline确认预测前没有额外做归一化训练集正常测试集偏色BatchNorm 的统计量偏移确认model.eval()被调用或替换为 InstanceNorm输出有棋盘格伪影转置卷积重叠把 ConvTranspose2d 换成 Upsample Conv2d高频细节全是噪点学习率太大从 1e-4 降到 3e-5生成图整体亮度偏低L1/L2 与数据归一化范围不匹配把目标数据 Normalize 到与模型输出相同的范围训练很慢GPU 利用率低DataLoader 的 num_workers 太小调大 num_workers确认图像解码是瓶颈“棋盘格伪影”是转置卷积自己叠加出来的高频噪声原因是卷积核滑动时产生重叠区域。如果发现生成的图片上有一块一块的网格纹理把上采样策略改成“先插值放大再做普通卷积”self.up nn.Upsample(scale_factor2, modebilinear, align_cornersFalse) self.up_conv nn.Conv2d(features[-1], features[-1] // 2, kernel_size3, padding1)这里features[-1]是 512经过up_conv后变成 256再和跳跃连接拼接回 512形状与原来用转置卷积一致。这个改动可以显著减少棋盘格代价是上采样过程不能学习空间位置信息实际效果通常更好。4.4 想进一步改进Unet模型先动哪两层如果目标是提升图像质量而不是追求模型复杂度优先看两个地方。一是把瓶颈层加厚。U-Net 的 bottleneck 只有一个 DoubleConv分辨率低但通道多负责压缩全局语义。把这里的通道从 1024 降到 512再接一个残差块可以减少参数又保持能力。二是加注意力门控class AttentionGate(nn.Module): def __init__(self, f_g, f_l): super().__init__() self.w_g nn.Conv2d(f_g, f_l, kernel_size1) self.w_x nn.Conv2d(f_l, f_l, kernel_size1) self.psi nn.Conv2d(f_l, 1, kernel_size1) self.sigmoid nn.Sigmoid() def forward(self, g, x): x1 self.w_x(x) g1 self.w_g(g) attn self.sigmoid(self.psi(F.relu(g1 x1))) return x * attn在解码器的每个层级用上采样后的特征图和跳跃连接的特征计算注意力图再让跳跃特征乘以注意力图。这样模型可以学会忽略无关的高频噪声聚焦与任务相关的结构区域。这是标准 Attention U-Net 的做法也是众多 U-Net 改进里最容易叠加进现有代码的模块之一。调这些之前别忘了先确认数据集够不够大、是否做了合理划分否则模型改进带来的收益很可能被过拟合覆盖。5. 用验证集和未见过的测试对检验你的Unet是不是真的学会了5.1 为什么不看训练集效果训练集 loss 降到 0只能说明网络死记住了这些样本的图像对应关系。特别是训练集图片只有几百张、背景非常相似时模型非常容易“背题”。判断是否真正学会必须在一个独立划分的验证集上评估。建议把数据按 8:1:1 分成训练、验证、测试。验证集参与调参比如决定什么时候停测试集只在所有超参数定完之后跑一遍用于评估最终效果。随机划分时要固定随机种子否则每次跑脚本数据分布都不一样改进没有可比性。import random import shutil from pathlib import Path random.seed(42) paths sorted(Path(./data/input).glob(*.png)) random.shuffle(paths) n len(paths) for idx, p in enumerate(paths): if idx int(n * 0.8): split train elif idx int(n * 0.9): split val else: split test target_file Path(./data/target) / p.name save_dir Path(f./data/{split}/input) / p.name target_dir Path(f./data/{split}/target) / p.name save_dir.parent.mkdir(parentsTrue, exist_okTrue) target_dir.parent.mkdir(parentsTrue, exist_okTrue) shutil.copy(p, save_dir) shutil.copy(target_file, target_dir)5.2 用PSNR和SSIM量化生成质量肉眼看图比较主观需要一个可重复的量化指标。PSNR 衡量像素误差SSIM 衡量结构相似度。两个指标一起看既能发现模糊问题也能发现结构扭曲问题。import torch from skimage.metrics import peak_signal_noise_ratio, structural_similarity def evaluate(model, loader, device): model.eval() psnr_total 0.0 ssim_total 0.0 count 0 with torch.no_grad(): for x, y in loader: pred model(x.to(device)).cpu() for p, t in zip(pred, y): p (p.numpy().transpose(1, 2, 0) 1) / 2 t (t.numpy().transpose(1, 2, 0) 1) / 2 psnr_total peak_signal_noise_ratio(p, t, data_range1) ssim_total structural_similarity(p, t, channel_axis2, data_range1) count 1 return psnr_total / count, ssim_total / count调用前记得model.eval()这样 BatchNorm 会用训练阶段缓存好的全局统计量而不是测试 batch 的统计量。PSNR 在背景平坦的图像上会虚高所以还要配合 SSIM 以及人眼观察。一个训练正常的模型在测试集上的 PSNR 通常会比训练集低 1 到 3dB如果差值超过 5dB说明已经过拟合需要增强数据或降低模型容量。5.3 最后做一个没见过的图像测试验证集和训练集同源仍然存在分布相同的可能。最后一步是找几张不在数据目录里的真实图片最好是场景、光线都和训练集不同的照片缩放裁剪到 256×256 后直接前向预测import torchvision.transforms as T from PIL import Image img Image.open(fresh_photo.jpg).convert(RGB) img T.Resize((256, 256))(img) img T.ToTensor()(img) img T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])(img).unsqueeze(0) model.eval() with torch.no_grad(): out model(img) out (out.squeeze(0).cpu().numpy().transpose(1, 2, 0) 1) / 2 Image.fromarray((out * 255).astype(uint8)).save(fresh_output.png)如果这个“新鲜”输入仍然输出合理结果说明 U-Net 学到的是可泛化的图像到图像映射而不是死记硬背。此时再回看训练日志记录下来的最优 epoch、损失函数配置、数据增强方式就是这次手写 U-Net 入门最有价值的产出后面做任何变体都可以拿它作为基准。本文还有配套的精品资源点击获取