Pytorch-UNet 图像语义分割指南：3 条命令训出可用的 U-Net 模型

Pytorch-UNet 图像语义分割指南3 条命令训出可用的 U-Net 模型【免费下载链接】Pytorch-UNetPyTorch implementation of the U-Net for image semantic segmentation with high quality images项目地址: https://gitcode.com/gh_mirrors/py/Pytorch-UNetPytorch-UNet 是一个开箱即用的 U-Net 图像语义分割 PyTorch 实现能把你的标注数据训成可用的模型再输出像素级掩码。如果你手头有一批带标注的图像想自己训一个分割模型最容易卡住的第一步是数据怎么摆、怎么确认结果在变好。这个项目把答案压缩成了 3 条命令——装依赖、下数据、开训官方记录中它在 5 千张图上从零训练在 10 万张以上测试图上拿到 0.988 的 Dice 系数。 Pytorch-UNet 能做什么一个跑得通的 U-Net 分割实现它按 U-Net 论文搭好结构把数据加载、训练、预测、评估、日志五件事封装成几个可直接执行的脚本。适合两类人想快速跑起一个分割实验的以及想读明白分割模型原理的。维度Pytorch-UNet自己手写 U-Net全功能分割框架如 MMSegmentation上手3 条命令出 checkpoint模型、损失、数据加载都要自己写配置项多学习曲线陡代码量网络定义 48 行看自己数千行调参方式命令行参数改源码配置文件适合场景快速实验、二次开发学习原理大型工程项目它替你做了三件事网络完整定义只有 48 行训练与预测两个脚本各一遍就能读完整个流程没有黑盒内置混合精度训练、Wandb 实时日志显存溢出时自动开启梯度检查点重试训练常见翻车点都提前处理过二分类与多分类分割之间只差一个--classes参数同一套代码能覆盖人像、医学、场景分割。⚡ 3 条命令跑通你的第一次训练假设 PyTorch 1.13 已按官网说明装好首次运行只需要下面 3 条命令# 安装其余 5 个依赖torch 本身需要单独安装 pip install -r requirements.txt # 下载 Carvana 汽车分割数据集约 5k 张高清图及掩码 bash scripts/download_data.sh # 混合精度训练每轮结束自动保存 checkpoint_epoch*.pth python train.py --ampscripts/download_data.sh 首次运行会向你要一个 Kaggle 用户名和 API key并替你写好到~/.kaggle/kaggle.json。训练启动后控制台会打印一条 wandb 链接打开就能看到损失曲线与预测掩码对比。结果在哪里看训练结束时checkpoints/下出现checkpoint_epoch1.pth说明第一次训练已经跑通。 核心能力拆解4 个模块各司其职模型定义3 个参数决定网络形状unet/unet_model.py 里是一个标准编码器-解码器结构4 层最大池化下采样把特征图越压越小4 层上采样再恢复每次上采样都把同尺度下采样侧的跳跃特征拼回来恢复边缘时不丢细节。搭一个网络只需 3 个参数model UNet(n_channels3, n_classesargs.classes, bilinearargs.bilinear)参数含义默认n_channels输入通道RGB 图用 3、灰度图用 13n_classes每像素输出几路概率二分类走 sigmoid、多类走 argmax2bilinearTrue 用双线性上采样省显存False 用转置卷积False显存特别紧张时调用model.use_checkpointing()前向逐层重算用计算换内存。训练循环双损失与自动调度train.py 里三行关键代码optimizer optim.RMSprop(model.parameters(), lrlearning_rate, ...) scheduler optim.lr_scheduler.ReduceLROnPlateau(optimizer, max, patience5) criterion nn.CrossEntropyLoss() if model.n_classes 1 else nn.BCEWithLogitsLoss()损失是「分类损失 Dice 损失」相加前者管每个像素分得对不对后者直接优化预测与真值掩码的重叠率和分割质量最直接相关验证 Dice 连续 5 次不涨就自动降学习率patience5显存溢出有兜底捕获torch.cuda.OutOfMemoryError后自动开启 checkpointing 重试日志同时提示开--amp。数据加载2 个目录和一套预处理约定utils/data_loading.py 定义了唯一的数据约定图像放data/imgs/、掩码放data/masks/同名一一对应且目录内不许有子目录或多余文件——加载器整目录扫描多一个文件就可能匹配出错。BasicDataset读入文件后做三件事按scale倍缩放默认 0.5掩码用最近邻插值保住边界图像归一化到 0-1把掩码像素值映射成类别索引 0..N-1原始值记在mask_values里预测时靠它把掩码还原成图片。Carvana 数据掩码文件名带_mask后缀会先尝试专用CarvanaDataset失败自动退回BasicDataset不用关心实际用了哪个类。评估与日志Dice 系数一个公式utils/dice_score.py 的核心指标是 Dice 系数——2 倍交集除以两者之和越接近 1 越好inter 2 * (input * target).sum(dimsum_dim) sets_sum input.sum(dimsum_dim) target.sum(dimsum_dim) dice (inter epsilon) / (sets_sum epsilon)训练中每 5 个 batch 跑一轮验证Dice 值连同损失、权重/梯度直方图、预测掩码样本一起写进 wandb学习率调度器也是按这个值走步。 场景实战Carvana 数据集上的两个例子训练一个汽车掩码模型输入data/imgs/的 5 千张 RGB 图与data/masks/的黑白掩码跑完下载脚本即就位。操作python train.py --epochs 20 --batch-size 4 --learning-rate 1e-4 --amp --scale 0.5这行把默认值改成 20 轮、batch 4、学习率 1e-4图像保持 0.5 倍缩放省显存并开启混合精度。输出checkpoints/里每轮一个.pthwandb 面板上能看到验证 Dice 逐步爬升控制台每隔一段打印一行Validation Dice score: 0.9x。新照片进来掩码出去输入一张不在训练集里的同场景汽车照片例如car.jpg。操作python predict.py -i car.jpg -o mask_car.png --model checkpoints/checkpoint_epoch20.pth输出mask_car.png车辆区域为白、其余为黑。predict.py 的预测核心逻辑是前向、把输出插值回原图尺寸、再按阈值二值化--mask-threshold控制像素判白的阈值默认 0.5调高掩码会更保守。不想落盘就加--viz --no-save窗口里直接看对比。 避坑指南训练时最常见的 5 种报错现象原因处理torch.cuda.OutOfMemoryError图像过大或没开混合精度加--amp、调低--scale、减小--batch-size程序也会自动开 checkpointing 重试No input file found in data/imgs目录里没有顶层图像只有子目录图像平铺摆放不要嵌套目录Either no image or multiple images found for the ID目录里存在同名多文件每个名字只保留一张图与一张掩码wandb 运行记录被自动删除未登录账号创建的是匿名 run设置WANDB_API_KEY环境变量绑定自己的账号验证 Dice 长时间不涨scale过小丢失细节或轮数不够显存允许时设--scale 1同时调大--epochs延伸从 Carvana 到你自己的任务的 3 条路跳过训练先用预训练权重项目通过 torch.hub 放出 Carvana 训练好的权重torch.hub.load(milesial/Pytorch-UNet, unet_carvana, pretrainedTrue, scale0.5)一行拿到可用网络先把自己的预测管线跑起来再谈微调换成自己的多分类任务数据保持同名一一对应掩码改用像素值为 0/1/2/3 的灰度图训练时加--classes 4损失会自动切到 softmax 交叉熵人像或医学分割走同一条流程容器化运行仓库自带 Dockerfiledocker build即可在 CUDA 基础镜像上构建训练容器适合放在共享 GPU 服务器上跑。项目地址git clone https://gitcode.com/gh_mirrors/py/Pytorch-UNet。原始论文为 U-Net: Convolutional Networks for Biomedical Image Segmentation读代码前先看一眼它的结构图再对照 unet_parts.py 里的Down与Up跳跃连接的位置会很直观。【免费下载链接】Pytorch-UNetPyTorch implementation of the U-Net for image semantic segmentation with high quality images项目地址: https://gitcode.com/gh_mirrors/py/Pytorch-UNet创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考