简介面向深度学习与医学图像分割初学者的PyTorch U-Net完整实现包提供了从网络搭建、数据预处理到训练测试的一整套流程可直接用于训练自定义图像分割数据集。压缩包共27个文件其中py源码覆盖核心模块如net.py、train.py、data.py、test.py及评估脚本png为可视化结果md为说明文档整体约602KB结构清晰便于对照学习。目前已有773人学习下载尤其适合希望理解编码器-解码器、跳跃连接等关键机制并动手实践自定义分割任务的读者。资源不仅包含可直接运行的U-Net模型与配套工具脚本还展示了数据标注与增强的处理思路可辅助完成从数据集准备到分割效果评估的完整闭环。1. 为什么拿 PyTorch 自己搭 U-Net图像分割没你想的那么玄用 PyTorch 搭建 U-Net 这件事网上一搜一大把教程但真到了要训练自己的数据集那一步翻车率其实很高。我见过太多人下了开源代码跑通了示例图换了自己的图片之后 loss 死活不降或者输出的 mask 全黑。问题大多不在网络结构而在于数据处理和训练细节。这份 pytorch-UNet 项目正好把训练、测试、评估的完整链路都串起来了代码量不大适合做图像分割的入门骨架。这篇文章我会按「网络结构 → 数据集制作 → 训练调参 → 避坑 → 验证部署」的顺序把它拆开讲适合正在做医学影像分割、卫星图地物提取或者工业质检的开发者。2. 拆解 U-Net 结构先从 net.py 看懂编码器、解码器和跳跃连接U-Net 的网络结构图网上到处都是对称的 U 型、左边收缩右边扩张但图和代码是两回事。真正让自己能改、能调还是要逐行读 net.py。这个文件里的实现是典型的原始 U-Net双卷积块 下采样 转置卷积上采样 跳跃连接。下面从最小单元开始讲。2.1 双卷积块是 U-Net 的最小单元U-Net 的基本构件不是单个卷积层而是「两次卷积 激活 归一化」的组合。原论文里没有 BatchNorm但这个项目的实现加了实际训练时 BN 对收敛帮助很大。import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super(DoubleConv, self).__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x)这里的 kernel_size3、padding1 是个关键选择它保证卷积前后特征图的宽高不变。如果不加 padding每次卷积尺寸都会变小下采样路径的尺寸对齐会变得很麻烦。BatchNorm 放在卷积和 ReLU 之间作用是对每个 batch 内的特征做归一化让激活值分布稳定。ReLU 用inplaceTrue节省显存在训练深网络时积少成多。第二个卷积的输出通道保持 out_ch 不变也就是说 DoubleConv 只做一次通道数改变空间尺寸全程不变。2.2 编码器下采样让通道数翻倍、尺寸减半编码器路径做的事情很规律先 MaxPool 把特征图缩小一半再进入 DoubleConv 把通道数翻倍。项目里这个操作被封装成了 Down 模块。class Down(nn.Module): def __init__(self, in_ch, out_ch): super(Down, self).__init__() self.mpconv nn.Sequential( nn.MaxPool2d(2), DoubleConv(in_ch, out_ch) ) def forward(self, x): return self.mpconv(x)nn.MaxPool2d(2)把 H 和 W 各缩小一半池化窗口是 2×2步长默认等于窗口大小。每经过一次 Down特征图的尺寸减半、通道数翻倍。在原论文里这个通道变化序列是 64 → 128 → 256 → 512 → 1024但这个项目在最后一个 Down 里做了改动用 512 → 512 而不是 512 → 1024这样做的直接好处是参数量减少对显存更友好。如果你的数据量不大这个改动反而能抑制过拟合。我在自己项目里试过原版 1024 的配置显存占用差了将近一倍但精度提升很有限。2.3 解码器转置卷积上采样拼接编码器特征解码器的核心是 Up 模块。它先用nn.ConvTranspose2d把特征图放大一倍然后把来自编码器的跳跃连接特征拼过来最后再过一次 DoubleConv。class Up(nn.Module): def __init__(self, in_ch, out_ch): super(Up, self).__init__() self.up nn.ConvTranspose2d(in_ch, in_ch // 2, 2, stride2) self.conv DoubleConv(in_ch, out_ch) def forward(self, x1, x2): x1 self.up(x1) diffY x2.size()[2] - x1.size()[2] diffX x2.size()[3] - x1.size()[3] x1 F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2]) x torch.cat([x2, x1], dim1) return self.conv(x)注意 forward 的参数x1 是来自上一层的上采样特征x2 是来自编码器同层的跳跃连接。F.pad是处理尺寸不整除情况的保险操作比如原图尺寸是奇数时下采样后特征图尺寸对不上直接拼接会报错。实际中我一般把输入图片统一 resize 成 16 的倍数这样 pad 分支基本不会触发但留着它能让网络更健壮。转置卷积的 kernel_size2、stride2 意味着上采样正好放大 2 倍。in_ch // 2是转置卷积的输出通道数这是为了在拼接后通道数回到期望值。拼接用的是torch.cat而不是相加这是 U-Net 的原始设计编码器的浅层特征保存着边缘和纹理信息解码器的深层特征语义更强拼接让两部分信息互补相当于给解码器开了一条「短接线」。2.4 forward 完整数据流从输入到输出的逐层走向把前面几个模块串起来就是完整的 UNet 类。class UNet(nn.Module): def __init__(self, n_channels, n_classes): super(UNet, self).__init__() self.inc DoubleConv(n_channels, 64) self.down1 Down(64, 128) self.down2 Down(128, 256) self.down3 Down(256, 512) self.down4 Down(512, 512) self.up1 Up(1024, 256) self.up2 Up(512, 128) self.up3 Up(256, 64) self.up4 Up(128, 64) self.outc nn.Conv2d(64, n_classes, 1) def forward(self, x): x1 self.inc(x) x2 self.down1(x1) x3 self.down2(x2) x4 self.down3(x3) x5 self.down4(x4) x self.up1(x5, x4) x self.up2(x, x3) x self.up3(x, x2) x self.up4(x, x1) return self.outc(x)forward 里的数据流很清晰x1 到 x5 是编码器的逐层输出up1 的输入是 x5 和 x4输出通道 256。之后每一层 Up 接收的通道数都是上一层的输出加上跳跃连接的通道。最后一个 Up 输出 64 通道经过 1×1 卷积nn.Conv2d(64, n_classes, 1)变成类别数。两个构造参数需要重点理解。n_channels是输入图像的通道数RGB 图填 3灰度图填 1。n_classes是分割类别数二分类任务填 1多分类填实际类别数。最后的 1×1 卷积相当于是把 64 通道的特征映射到类别空间二分类时输出单通道后面接 sigmoid。3. 数据集准备把原始图片整理成 JPEGImages 和 SegmentationClassU-Net 的数据集组织方式沿用了 VOC 分割的目录结构原始图片放 JPEGImages标注 mask 放 SegmentationClass。这个项目里还带了一个 make_mask_data.py 脚本专门处理「把标注转成模型能读的 mask」这件事。很多人拿到项目第一步就卡在这里因为网上能找到的示例数据都是处理好的自己的标注却五花八门。3.1 目录约定VOC 格式的命名与类别映射项目根目录下有两个关键文件夹JPEGImages 和 SegmentationClass。前者放原始图片后者放对应的分割标签图。命名必须一一对应举个例子JPEGImages/001.jpg对应的 mask 就是SegmentationClass/001.png。类别映射是整个数据准备里最容易出问题的环节。VOC 格式的 mask 图里背景像素是黑色RGB 值为 0,0,0目标区域用纯色块填充比如红色代表类别 1绿色代表类别 2。模型训练时并不直接读这些 RGB 值而是要把它们转换成单通道索引图每个像素存一个整数0 表示背景1 表示第一个类别2 表示第二个类别。我一般会用下面这个函数来检查自己的 mask 到底有几个类别import numpy as np from PIL import Image mask Image.open(SegmentationClass/001.png) arr np.array(mask) print(mask shape:, arr.shape) print(unique values:, np.unique(arr))关键要看输出。如果 mask 是单通道灰度图unique values 应该是一组小整数比如 [0, 1, 2]。如果 mask 是三通道 RGBunique values 会是一大串 0 到 255 的值那就需要先量化成索引。如果出现 255说明标注里用了白色表示目标但模型期望的是 0 和 1这个不处理训练出来的结果必然不对。3.2 make_mask_data.py 的 mask 制作逻辑这个脚本做的事情就是把标注图从「人眼看的格式」转成「模型吃的格式」。核心逻辑是像素级颜色映射import os import numpy as np from PIL import Image color_to_class { (0, 0, 0): 0, # 背景 (255, 0, 0): 1, # 类别1红色 (0, 255, 0): 2, # 类别2绿色 } def make_mask(src_dir, dst_dir): os.makedirs(dst_dir, exist_okTrue) for name in os.listdir(src_dir): if not name.lower().endswith((.png, .jpg, .jpeg)): continue rgb Image.open(os.path.join(src_dir, name)).convert(RGB) arr np.array(rgb) h, w, _ arr.shape mask np.zeros((h, w), dtypenp.uint8) for color, cls in color_to_class.items(): match (arr[..., 0] color[0]) \ (arr[..., 1] color[1]) \ (arr[..., 2] color[2]) mask[match] cls Image.fromarray(mask).save(os.path.join(dst_dir, os.path.splitext(name)[0] .png))这段代码的逻辑很直白遍历每一张原始图片把 RGB 像素值和预设的颜色字典比对命中的像素位置写入对应的类别编号。最后保存成单通道 PNGPNG 格式支持单通道灰度图不会丢信息。注意dtypenp.uint8不能省如果不指定默认生成的数组是 int64PIL 保存时会报错或者写出 16 位图。实际项目里 labels 往往不是纯色块尤其医学影像里很多标注工具生成的是带边缘羽化的 PNG。这种情况我一般会先把像素 RGB 值转成 HSV再按色相区间归类或者直接用标注工具导出的索引 PNG而不是在脚本里做颜色猜测。make_mask_data.py 适用于「颜色规范、边界清晰」的标注如果你的标注来源复杂先把颜色统一再做映射。3.3 数据增强与归一化img 和 mask 必须同步变换数据集类里最容易踩的坑是图片做了旋转翻转mask 也跟着变但很多人写增强时只对 img 做了变换mask 没动。结果训练时模型看到的是「错位的标签」loss 高到爆炸还找不到原因。import os import random from PIL import Image from torch.utils.data import Dataset class VOCDataset(Dataset): def __init__(self, img_dir, mask_dir, size(512, 512)): self.img_dir img_dir self.mask_dir mask_dir self.names [n for n in os.listdir(img_dir) if n.lower().endswith((.png, .jpg, .jpeg))] self.size size def __getitem__(self, idx): name self.names[idx] img Image.open(os.path.join(self.img_dir, name)).convert(RGB) mask Image.open(os.path.join(self.mask_dir, os.path.splitext(name)[0] .png)) img img.resize(self.size, Image.BILINEAR) mask mask.resize(self.size, Image.NEAREST) if random.random() 0.5: img img.transpose(Image.FLIP_LEFT_RIGHT) mask mask.transpose(Image.FLIP_LEFT_RIGHT) if random.random() 0.5: img img.transpose(Image.FLIP_TOP_BOTTOM) mask mask.transpose(Image.FLIP_TOP_BOTTOM) img np.array(img, dtypenp.float32) / 255.0 mask np.array(mask, dtypenp.int64) img img.transpose(2, 0, 1) # HWC - CHW return torch.from_numpy(img.copy()), torch.from_numpy(mask.copy())注意两个 resize 的插值方式不同。图像用 BILINEAR 双线性插值保留平滑的边缘mask 必须用 NEAREST 最近邻插值否则类别边界会出现「新类别」——比如 0 和 1 之间插值出 0.5模型就懵了。随机翻转时 img 和 mask 要执行同一种变换上面代码用的两个独立 if 块其实有问题应该用一个随机种子或者同时翻转。这是我实际写代码时踩过的坑正确写法是用seed random.random()决定是否翻转img 和 mask 用同一个 seed。归一化方面图像除以 255 后取值范围是 0 到 1足够用了。ImageNet 的 mean/std 归一化在 U-Net 上不一定更好我试过几次对分割任务提升不明显反而多一步换算。mask 保持 0、1、2 这样的整数不做归一化交叉熵损失要求的是索引值而不是 one-hot 概率。3.4 数据量不足切 patch 和离线增强是补救手段医学图像分割经常遇到「标注图只有三五十张」的情况。U-Net 虽然比大模型的参数少但没有预训练权重时30 张图根本喂不饱。常见做法是切 patch把 1024×1024 的大图切成 256×256 的 patch相邻 patch 之间可以设 50% 重叠相当于把数据量翻了十几倍。训练时额外加在线增强包括 90 度旋转、缩放、亮度抖动、高斯噪声。如果还是不够再考虑用更大的切图步长生成更多 patch。但要注意切 patch 后要让训练集和验证集来自不同的原图否则模型会「记住」重叠区域验证指标虚高。4. 训练自己的数据集train.py 的参数设置和损失函数选型数据集就绪之后训练脚本就是整个项目的引擎。train.py 里做了数据加载、模型初始化、前向传播、损失计算、反向传播和模型保存。这个环节的选型直接决定模型能不能收敛。4.1 训练入口与数据加载先确认 PyTorch 环境开始训练之前先保证基础环境是干净的。我的习惯是 conda 新建一个独立环境避免系统 Python 里一堆包的版本冲突。如果你还没搭好 PyTorch 环境下面是常见做法conda create -n unet python3.8 conda activate unet pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118装完后可以用一行命令验证 CUDA 是否可用import torch print(torch.cuda.is_available(), torch.cuda.get_device_name(0))这里有个常见问题很多人装的是 CPU 版 PyTorch训练时才发现torch.cuda.is_available()返回 False白白跑了一晚上。所以我一般会在 train.py 开头强制检查设备状态而不是靠运气。数据加载部分用 PyTorch 的DataLoader核心参数是 batch_size、shuffle 和 num_workersfrom torch.utils.data import DataLoader train_dataset VOCDataset(JPEGImages, SegmentationClass, size(512, 512)) train_loader DataLoader(train_dataset, batch_size4, shuffleTrue, num_workers2) model UNet(n_channels3, n_classes1) device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device)num_workers是数据加载的子进程数Windows 上设成 0 最稳Linux 可以设成 2 或 4。shuffleTrue 必须在每个 epoch 打乱数据顺序否则模型会学到数据的排列顺序而不是语义特征。4.2 损失函数BCE 和 Dice Loss 怎么选二分类分割任务里这个项目用的是nn.BCEWithLogitsLoss。它的特点是内部先做了 sigmoid 再算交叉熵数值上比「自己算 sigmoid BCE」更稳定因为避免了 log(0) 的情况。多分类任务则用nn.CrossEntropyLoss它内部做了 softmax输入的是网络最后的原始 logits。医学分割场景下我更推荐把 Dice Loss 也试一下。Dice 系数衡量的是预测和真实 mask 的重叠率它对前景占比很小的图特别友好因为交叉熵在小目标上容易被背景淹没。常用的组合是bce dice两者加权相加import torch.nn.functional as F def bce_dice_loss(pred, target, alpha0.5): bce F.binary_cross_entropy_with_logits(pred, target.float()) pred_prob torch.sigmoid(pred) smooth 1.0 dice 1 - (2 * (pred_prob * target).sum() smooth) / \ (pred_prob.sum() target.sum() smooth) return alpha * bce (1 - alpha) * dice这里target需要是 float 类型且形状和 pred 一致。smooth 参数是为了防止分子分母都是 0 时除零错误一般设 1 就够了。alpha 控制两个损失的比例常见做法是 0.5 对半分。如果你的数据集正负样本比例严重失衡比如病灶只占整张图的 5%建议把 alpha 调低到 0.3让 Dice 占主导。4.3 优化器、学习率、batch size 和权重初始化优化器这块没什么悬念Adam 是 U-Net 训练的首选。它的自适应学习率让新手不用频繁调参收敛速度也快。超参数推荐值备注optimizerAdam默认 betas(0.9, 0.999)初始学习率1e-3如果 loss 震荡降到 1e-4batch_size4~8取决于显存512×512 输入通常 4 就够epochs50~200小数据集关注验证集不要盲目堆 epoch输入尺寸512×512需要能被 16 整除import torch.optim as optim optimizer optim.Adam(model.parameters(), lr1e-3) criterion nn.BCEWithLogitsLoss() for epoch in range(100): model.train() for imgs, masks in train_loader: imgs, masks imgs.to(device), masks.to(device).float() preds model(imgs) loss criterion(preds, masks.unsqueeze(1)) optimizer.zero_grad() loss.backward() optimizer.step() print(fEpoch {epoch}, Loss: {loss.item():.4f})masks.unsqueeze(1)把 [B, H, W] 变成 [B, 1, H, W]因为网络输出是单通道必须对齐维度。如果忘记做这一步PyTorch 会报维度不匹配的错误这是新手最常见的问题。optimizer.zero_grad()每次迭代都要调用否则梯度会累加。学习率方面小数据集上我习惯用ReduceLROnPlateau来降低学习率验证集 loss 连续 5 个 epoch 不降就除以 10。权重初始化也值得注意。PyTorch 的 Conv2d 默认初始化是 Kaiming 均匀分布U-Net 结构比较深直接训练基本没问题。但如果你的数据量很少建议载入 ImageNet 上预训练的 encoder 部分再微调解码器。5. 避坑指南U-Net 训练自己的数据集常见的 5 个翻车现场训练 U-Net 的过程本质上是在和数据格式、维度、显存作斗争。下面这几条都是我多次见过的实际问题每一条都是「现象 → 原因 → 解决」的记录。5.1 mask 是三通道 RGB训练时 loss 死活不降现象训练时 loss 维持在 0.69 左右不动输出图全黑或者全白。原因mask 是从标注工具直接导出的 RGB 三通道 PNG不是单通道索引图。模型输出是单通道但计算损失时 mask 是 [B, 3, H, W]两个张量对不上模型只能在三个通道里「猜」。解决用 3.2 节的make_mask_data.py把 RGB mask 转成单通道索引图。转换后检查一下np.unique(mask)的输出确保只有 0、1、2 这种小整数。5.2 mask 像素值是 0 和 255不是 0 和 1现象训练能跑loss 也在下降但验证集准确率低得离谱。原因很多标注工具把目标区域标成纯白255背景标成纯黑0。BCEWithLogitsLoss 期望的目标是 0 到 1 之间的概率用 255 和 0 去算交叉熵相当于用一个被放大的错误值去更新梯度。解决训练脚本里对 mask 做一次归一化mask (mask 0).float()把大于 0 的值全部变成 1。注意这一步在 DataLoader 里做而不是在数据预处理阶段做。5.3 输入尺寸不是 16 的倍数训练时跳跃连接报尺寸错误现象报错信息形如RuntimeError: Sizes of tensors must match except in dimension 1。原因U-Net 有 4 次下采样特征图尺寸每层减半。如果输入尺寸不是 16 的倍数经过 4 次池化后不同分支的特征图尺寸会出现 1 像素的差异拼接时直接报错。解决所有图片统一 resize 到 16 的倍数比如 256、320、512。如果你不想 resize也可以像我一样在 Up 模块里保留F.pad补齐但尽量用统一尺寸让数据分布更稳定。5.4 CUDA 显存溢出batch_size 调到 1 还是不够现象CUDA out of memory程序直接崩。原因输入图分辨率太高比如原始 CT 图是 1024×1024即使 batch_size1前向传播时中间特征图的显存开销也很大。解决两个方案。一是把输入尺寸改成 512×512 或 256×256二是改用 patch 训练把大图切成 256×256 的小块。医学影像场景里我一般用后者因为下采样后的特征图能保留更多细节。推理时再用 overlap-tile 策略拼接回原尺寸。5.5 训练集只有几十张图模型过拟合严重现象训练集 loss 降到 0.01验证集 Dice 只有 0.5训练曲线前后差距巨大。原因模型参数几百万数据量太少网络把训练集的噪声细节全背下来了。解决数据增强 减少模型容量。U-Net 的通道数可以从 64 起步砍成 32参数量直接减少 4 倍。还可以加 Dropout或者用带预训练的 encoder 做迁移学习。如果标注成本可以接受优先找更多数据这是最踏实的解法。6. 验证与部署test.py 结果可视化以及模型转 ONNX 的实用技巧训练完成后验证和部署是真正检验模型价值的两步。test.py 加载训练好的权重对新图片做前向推理输出分割 mask。评估脚本 get_evaluation.py 会计算 Dice 和 IoU 指标值得在每次训练后都跑一遍。model UNet(n_channels3, n_classes1) model.load_state_dict(torch.load(model.pt, map_locationcpu)) model.eval() with torch.no_grad(): pred model(img.unsqueeze(0)) pred torch.sigmoid(pred).squeeze().numpy() pred_binary (pred 0.5).astype(np.uint8) * 255img.unsqueeze(0)是为了把 [3, H, W] 变成 [1, 3, H, W] 的 batch 维度。model.eval()必须显式调用它会关闭 Dropout 和 BatchNorm 的统计更新否则同一张图推理两次结果会不一样。阈值 0.5 是二分类的默认选择如果你的类别严重不平衡需要根据验证集的 precision-recall 曲线调整。模型导出 ONNX 是部署到实际工程里的关键一步。常见的做法是导出固定尺寸的 ONNX再验证输出一致性dummy_input torch.randn(1, 3, 512, 512) torch.onnx.export( model, dummy_input, unet.onnx, opset_version11, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}} )这里只把 batch 维设为动态输入分辨率保持固定因为很多部署框架对动态分辨率支持不稳定。导出前模型必须处于 eval 状态否则 BatchNorm 会被封装进计算图。之后用 onnxruntime 加载和 PyTorch 推理结果对比最大误差一般应该小于 1e-5如果差异明显多半是预处理方式不一致。从那以后我每次训练完都会强制走一遍「test 单张图 → 算 Dice → 导出 ONNX → 对比一致性」这条流程而不是只看训练 loss。没有这最后一步模型在验证集上再漂亮部署时也可能翻车。希望帮到你。本文还有配套的精品资源点击获取