简介这是一份基于PyTorch实现Unet多类别语义分割的实战资源面向具备一定深度学习基础、希望掌握图像分割模型训练与调优的开发者。资源共46个文件以Python源码为核心包含19个py脚本与24个pyc编译文件另有2个txt配置说明和1个json参数文件整体仅69KB结构紧凑便于直接阅读与修改。已有15245人学习这份资源属于同类教程中热度较高的实践型内容。作者围绕Encoder-Decoder结构、跳跃连接、多类别输出通道设计等关键点组织代码配套提供数据加载、自定义变换、训练评估、损失函数、学习率调度、指标计算等模块并附有数据集划分与类别权重计算脚本基本覆盖了从数据预处理到模型推理的完整流程。通过阅读源码与对应博客读者可掌握Unet在自定义多类别数据集上的迁移方法、训练流程设计及调优技巧适合用于医学影像、遥感图像等场景的入门与项目参考。1. 项目概述为什么选择Pytorch配合Unet做多类别分割先说结论如果你手里有一批自己的图片数据想按像素把图中不同物体分出来比如道路、建筑、植被、水体这类地物目标那“Pytorch Unet 多类别数据集”这套组合是目前开源社区里最省心、最不容易把自己绕晕的路线。这个标题之所以有这么多人在搜是因为它几乎覆盖了从零入门语义分割的所有关键环节框架选型、模型结构、数据组织、训练调参、结果评估。我自己最早接触语义分割时也纠结过用TensorFlow还是Pytorch后来彻底转向Pytorch原因很朴素调试方便报错信息看得懂断点能直接打在张量运算那一行上。配合Unet这种编码器-解码器结构哪怕数据集只有几百张图也能训出一个效果不错的多类别分割模型。这篇文章不会跟你聊太虚的理论而是把我实际跑通整个流程的步骤、参数、踩坑记录都摆出来你照着操作就能在自己的多类别数据集上复现。文章适合三类人看刚入门语义分割的学生、需要用自己数据做分割实验的工程师、以及想快速验证Unet效果的产品人员。2. 整体设计思路与方案选型2.1 为什么Unet依然是多类别分割的首选基线Unet之所以经典核心在于它同时保住了“细节”和“语义”。下采样路径不断缩小特征图尺寸让模型能看到更大的感受野上采样路径则把高层的语义信息逐步还原到原图分辨率。中间那一圈跳连接把下采样时各层的位置细节直接拼到上采样路径上相当于给模型开了一条“记忆通道”小目标边缘不容易丢。对于多类别分割任务比如五类甚至十类地物Unet每一条跳连接都在帮助模型区分“这里是边界还是内部”。我自己做过对比实验在同样的数据集上把Unet换成PSPNet发现小目标类别的交并比下降明显。原因不复杂PSPNet把重点放在全局池化上对小目标的敏感度反而不如Unet这种逐层传递细节的结构。所以如果你的数据里存在较多小目标或细长条目标Unet是最稳的起点。2.2 Pytorch生态里的三个关键选择第一框架版本跟进。建议用Pytorch 2.x系列如果你需要GPU训练记得提前确认CUDA、cuDNN和显卡驱动的匹配关系。以2024年之后的环境为例Pytorch 2.0以上版本对自动混合精度训练的支持更完善显存占用更友好。第二模型实现方式。可以直接从网上找Unet的Pytorch实现也可以自己按论文结构手写。实话说手写一遍Unet比复制十遍别人的代码都有用结构细节会刻在你脑子里。第三预训练编码器。如果你用torchvision里的ResNet作为Unet的骨干加载ImageNet预训练权重训练收敛速度会快不少尤其是当你的数据集规模不大的时候。选型逻辑很直接小数据集靠预训练权重大数据集靠模型容量。如果你的数据只有两三百张建议选择ResNet34做编码器如果数据到了一千张以上可以尝试ResNet50或者直接换EfficientNet。这个不是铁律而是我在不同规模数据上反复试出来的经验。3. 数据准备多类别数据集的整理与预处理3.1 目录结构与标注格式的统一很多人在模型跑不起来时才发现问题出在数据上而不是代码上。多类别语义分割的数据集标准的组织方式是这样dataset/ ├── images/ │ ├── img_001.jpg │ ├── img_002.jpg │ └── ... └── masks/ ├── img_001.png ├── img_002.png └── ...图片格式一般用jpg或png都行但掩膜mask必须是png而且是单通道的灰度图或调色板模式。为什么必须png因为jpg是有损压缩会导致标注类别边缘出现伪色模型会学到错误信息。特别注意掩膜每个像素点的数值背景为0第一个类别为1第二个类别为2以此类推。如果你用Labelme这类工具标注导出的掩膜是调色板模式需要用代码转成类别索引不然训练时一算损失就是一片NaN。3.2 数据增强与样本均衡技巧多类别分割里最头疼的问题就是类别不平衡。比如一栋建筑物可能只占画面面积的5%背景却占了70%。如果直接训练模型会倾向把所有像素都预测成背景。我的做法是对每个类别统计像素占比然后给损失函数里的每个类别分配权重权重和像素占比成反比。数据增强我用的是albumentations库比torchvision的transform灵活得多。我常用的一套增强组合包括水平翻转、垂直翻转、随机旋转90度、随机裁剪和亮度对比度调整。这里有个细节对图像做翻转和旋转时掩膜必须做同样的变换albumentations保证了这一点。增强不是越多越好如果你的类别是建筑物这类刚性目标翻转和旋转没问题如果你处理的是文本行这类有方向性的目标旋转90度会把标注搞乱。3.3 自定义Dataset类的关键代码写Dataset类时最容易犯的一个错误是忘记把掩膜里的类别索引压到从0开始连续分布。以下是我常用的代码import torch from torch.utils.data import Dataset from PIL import Image import os import numpy as np class SegmentationDataset(Dataset): def __init__(self, image_dir, mask_dir, transformNone, class_mappingNone): self.image_dir image_dir self.mask_dir mask_dir self.transform transform self.class_mapping class_mapping or {} self.images sorted([f for f in os.listdir(image_dir) if f.endswith((.jpg, .png))]) def __len__(self): return len(self.images) def __getitem__(self, idx): img_name self.images[idx] img_path os.path.join(self.image_dir, img_name) mask_path os.path.join(self.mask_dir, img_name.replace(.jpg, .png)) image np.array(Image.open(img_path).convert(RGB)) mask np.array(Image.open(mask_path)) # 不要convert(RGB)保持灰度 # 如果mask是调色板模式0-255任意值需要做类别映射 if self.class_mapping: mapped_mask np.zeros_like(mask) for old_id, new_id in self.class_mapping.items(): mapped_mask[mask old_id] new_id mask mapped_mask if self.transform: transformed self.transform(imageimage, maskmask) image transformed[image] mask transformed[mask] image_tensor torch.from_numpy(image).permute(2, 0, 1).float() / 255.0 mask_tensor torch.from_numpy(mask).long() return image_tensor, mask_tensor这段代码里有个非常重要的点掩膜转成tensor用.long()不要用.float()。因为后面损失函数CrossEntropyLoss期望输入是整数类别索引如果用float会报错或者产生错误结果。我在一开始就踩过这个坑整整调了一个晚上才发现是类型不匹配。4. 模型实现Unet结构与多类别输出适配4.1 Unet核心结构拆解我在这里不放完整的三百行Unet代码了因为网上开源实现非常多我建议你找一个star数高的仓库读一遍结构但有几个核心参数必须搞清楚。Unet整体分为编码器、瓶颈、解码器三段。编码器是若干个卷积块加下采样特征图尺寸减半通道数翻倍瓶颈在最底层解码器逐步上采样通道数减半并与对应的编码器特征图拼接。关键在于最后一层卷积的输出通道数必须等于你的类别数。举例如果是五类分割最后一层卷积输出通道数就设为5。每个通道对应一个类别的置信度分数。训练时用CrossEntropyLoss它对每个像素在通道维度上做softmax后计算损失。这部分不需要你自己写softmaxPytorch的CrossEntropyLoss内部已经包含了。4.2 多类别输出的通道设置与损失函数选择我遇到不少人在修改Unet时只改了模型最后一层的输出通道数但忽略了编码器预训练权重的加载方式。如果是自己写的Unet从头训练没问题如果你的编码器要加载预训练权重前几层的通道数必须和ImageNet预训练模型一致通常就是RGB三通道输入输出通道按骨干网络设定。损失函数方面多类别分割最常用的是CrossEntropyLoss加DiceLoss的组合。我实际测试下来纯用CrossEntropyLoss小目标类别的分割效果一般纯用DiceLoss训练初期损失波动剧烈。两者的加权和比较稳具体公式是loss 0.7 * ce_loss 0.3 * dice_loss这个比例可以根据你的数据调整。如果类别特别不均衡把dice_loss的权重调高到0.5甚至0.7。DiceLoss对不平衡不敏感它能直接优化类别区域的重合度。5. 训练配置与完整实操流程5.1 环境准备与关键参数设置环境方面推荐用Anaconda创建独立的环境。以Ubuntu系统为例常见组合是Python 3.10.11加Pytorch 2.8.0加CUDA 12.1这一套搭配在GTX 30系和40系显卡上表现稳定。Windows下的流程类似只是CUDA环境变量配置要格外小心。如果显卡显存只有6G建议输入图片尺寸用256x256批次大小设为4到8显存12G以上的话输入尺寸可以提高到512x512批次大小设8到16。我整理了训练阶段几个关键的超参数参考值参数推荐值备注输入尺寸256x256 / 512x512小显存用256大显存用512批次大小4 / 8 / 16视显存而定初始学习率1e-4 / 1e-3使用AdamW优化器学习率调整CosineAnnealingLR避免后期震荡训练轮数50 / 100看验证集指标早停优化器AdamW比Adam稳定性好5.2 训练循环中的关键代码训练循环本身并不复杂但有一些细节会让训练过程顺畅很多。比如启用自动混合精度用autocast和GradScaler显存能省下近一半训练速度还能提升。另外每个epoch保存一次checkpoint别只保存在最后一个epoch——训练中断是家常便饭有checkpoint才能续上继续训。一个简化但完整可跑的训练循环结构import torch from torch.cuda.amp import autocast, GradScaler scaler GradScaler() optimizer torch.optim.AdamW(model.parameters(), lr1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_maxepochs) for epoch in range(epochs): model.train() total_loss 0 for images, masks in train_loader: images, masks images.to(device), masks.to(device) optimizer.zero_grad() with autocast(): outputs model(images) loss ce_criterion(outputs, masks) 0.3 * dice_loss(outputs, masks) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() total_loss loss.item() scheduler.step() avg_loss total_loss / len(train_loader) print(fEpoch {epoch}/{epochs}, Loss: {avg_loss:.4f}) if epoch % 5 0: torch.save(model.state_dict(), fcheckpoints/unet_epoch_{epoch}.pth)注意dice_loss需要自己实现def dice_loss(pred, target, smooth1.0): pred torch.softmax(pred, dim1) target_onehot torch.nn.functional.one_hot(target, num_classesnum_classes).permute(0, 3, 1, 2).float() intersection (pred * target_onehot).sum(dim(2, 3)) union pred.sum(dim(2, 3)) target_onehot.sum(dim(2, 3)) dice (2 * intersection smooth) / (union smooth) return 1 - dice.mean()5.3 评估指标与结果可视化训练完成后不要只看loss要算每个类别的IoU和mIoU。mIoU是所有类别IoU的平均值这是语义分割最通用的评价指标。我需要提醒你计算IoU时要对每个类别单独统计不能直接拿混淆矩阵整体算。具体实现可以用sklearn的confusion_matrix辅助。from sklearn.metrics import confusion_matrix def compute_iou(pred_mask, true_mask, num_classes): pred_flat pred_mask.flatten() true_flat true_mask.flatten() cm confusion_matrix(true_flat, pred_flat, labelslist(range(num_classes))) intersection np.diag(cm) union cm.sum(axis0) cm.sum(axis1) - np.diag(cm) iou intersection / (union 1e-6) return iou可视化方面推荐把预测结果叠加到原图上用半透明色块显示不同类别。我觉得看分割效果比死磕指标更直观很多边界问题光看IoU是发现不了的。别只看训练集的预测效果一定要去验证集上抽几张图看边界质量——是不是出现锯齿状、是不是有空洞。6. 常见问题与排查技巧实录6.1 损失降至Nan的排查过程这是最多人碰到的坑。我遇到过一次损失降到NaN排查步骤是先检查掩膜里是否有超出类别数的值。比如类别数是5掩膜数值范围却是0到255这会让CrossEntropyLoss计算出NaN。然后检查数据归一化图片输入是否除以了255掩膜是否保持整数类型。最后检查学习率如果初始学习率设到0.1梯度爆炸也会导致NaN。我习惯把初始学习率控制在1e-4到1e-3配合AdamW很少再碰到NaN。6.2 显存不足的应对策略显存不足是个很现实的问题。我一开始用512x512输入、批次16直接爆显存。后来做了三件事换成分批训练、缩小输入尺寸、启用混合精度。假设你的显卡是8G显存推荐直接用256x256输入加批次8加混合精度这样训练速度反而可能比大尺寸低批次更稳定。另外在forward里加一句torch.cuda.empty_cache()也能清理一部分碎片显存但不要在每个step都调用会拖慢速度。6.3 模型训练不收敛或过拟合的调整训练不收敛先看损失曲线是震荡还是不下降。震荡说明学习率太高调低一个数量级不下降说明可能模型结构或数据出了问题。我遇到过一次模型输出恒为背景的情况检查发现掩膜数值没有对齐所有类别都被映射成了0。过拟合则表现为训练集损失低但验证集IoU不增这时加大数据增强强度、加Dropout、缩小模型容量都有效。我把这些高频问题整理成了一个速查表症状可能原因解决方案损失NaN掩膜值超出类别范围检查mask像素值确保0到num_classes-1显存不足输入尺寸过大/批次过大降输入尺寸、分批、开混合精度验证IoU不涨过拟合或学习率过大加增强、降学习率、早停预测全是背景类别索引不对齐检查数据集类映射逻辑边界粗糙跳连接被忽略检查模型是否真的用到了skip connection6.4 模型保存、加载与推理的完整流程训练完模型后保存方式我推荐只存state_dict不存整个模型因为后者在Pytorch版本升级后容易反序列化失败。加载模型后做推理要注意输入图片必须做和训练时一样的预处理resize到相同尺寸、归一化到0到1、转成张量、加batch维度。预测输出是一个形状为(1, num_classes, H, W)的张量用argmax(dim1)取每个像素的类别索引。最后转换成彩色图时准备一个调色板数组把类别索引映射成RGB颜色。import torch import numpy as np from PIL import Image import torchvision.transforms as transforms def inference_single_image(model, image_path, device, num_classes): transform transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) image Image.open(image_path).convert(RGB) input_tensor transform(image).unsqueeze(0).to(device) model.eval() with torch.no_grad(): output model(input_tensor) pred torch.argmax(output, dim1).squeeze(0).cpu().numpy() return pred推理后如果想保存成彩色分割图可以用这样一个简单的映射color_map np.array([ [0, 0, 0], # 背景 [255, 0, 0], # 类别1红色 [0, 255, 0], # 类别2绿色 [0, 0, 255], # 类别3蓝色 [255, 255, 0], # 类别4黄色 ], dtypenp.uint8) segmentation_rgb color_map[pred] Image.fromarray(segmentation_rgb).save(prediction.png)7. 多类别数据标注与类别映射的实战心得7.1 标注工具选择与格式转换常见坑多类别分割最关键的第一步其实是标注质量。我试过Labelme、EISeg、CVAT最终常用Labelme配合脚本转成Unet需要的png掩膜。Labelme导出的是json文件每个多边形对应一个label需要逐张解析json并把多边形填充成掩膜。这里有个非常隐蔽的坑json里标注的label名称和你最终想要的类别索引可能不是一回事一定要建立一个字典做映射。import json import numpy as np import cv2 import os def json_to_mask(json_path, height, width, label_map): with open(json_path, r, encodingutf-8) as f: data json.load(f) mask np.zeros((height, width), dtypenp.uint8) for shape in data[shapes]: label shape[label] points np.array(shape[points], dtypenp.int32) if label in label_map: cv2.fillPoly(mask, [points], label_map[label]) return mask这段函数把json里的每个多边形填充到掩膜上label_map里存的是比如{建筑: 1, 道路: 2, 植被: 3}这样的键值对。实际操作中我发现最花时间的不是写转换脚本而是清洗标注数据。比如相邻图片边缘处多边形没有贴合图像边界导致交界处出现一条无类别带状区域这一条区域会成为模型预测错误的高发区。处理方法是对掩膜做一个形态学闭运算把细小的空洞缝补上。7.2 数据集划分比例与验证集选择多类别分割的数据划分和普通分类不太一样。除了按文件数量比例划分还要考虑类别分布。我遇到过一种情况训练集里“水体”样本很多验证集里“水体”只出现在一张图的一角结果导致验证集水体IoU特别低模型并没有过拟合纯粹是验证集抽样偏差。比较稳妥的做法是按图像整体划分保证每一类在训练集和验证集都出现如果某个类别的图像特别少可以考虑将数据增强用到验证集上但这并不常规更合理的是用k-fold交叉验证。对数据量只有几百张的情况k-fold交叉验证是评估模型真实水平的有效方式。最后分享一个细节在训练过程中如果发现某几个类别的IoU一直偏低先别急着改模型结构回到标注数据里看看这些类别的标注质量是否有重叠标注、漏标、边界粗糙的情况。我之前处理遥感影像分割时“阴影”类别的IoU怎么都上不去后来仔细拉大图片比对发现标注员把很多本来就模糊的阴影边界标歪了模型学到的边界自然混乱。重新清洗了一批标注后这个类别的IoU直接提升了十几个点。数据质量决定了模型效果的上限这话在语义分割里体现得特别明显。本文还有配套的精品资源点击获取