企业官网建设流程全解析

Pytorch-UNet 图像分割从跑通到部署：新手完整指南

【免费下载链接】Pytorch-UNetPyTorch implementation of the U-Net for image semantic segmentation with high quality images项目地址: https://gitcode.com/gh_mirrors/py/Pytorch-UNet

分割模型最难啃的往往不是网络本身，而是环境没配好、数据读不进来、显存总溢出的那些琐碎环节。这篇围绕 Pytorch-UNet 图像分割训练展开，用项目自带的 Carvana 数据集做示例，几条短命令走通从冒烟测试到拿到可交付掩码的全过程。你只要会跑git clone和pip install就够了。

先跑通：最短路径拿到第一个验证结果

新手最容易犯的错误是一上来就精调参数。先让流水线转起来，看到第一个数字，后面调参心里才有底 🚀

git clone https://gitcode.com/gh_mirrors/py/Pytorch-UNet cd Pytorch-UNet pip install -r requirements.txt bash scripts/download_data.sh python train.py --epochs 1 --amp

第二条之后的命令都在仓库根目录执行。scripts/download_data.sh会交互式地要你的 Kaggle 用户名和 API key（Windows 用户改用scripts\download_data.bat），然后把约 5000 张带标注的汽车图像拉进data/。最后一条命令只跑 1 个 epoch，目的是验证环境、数据、训练三件事都没问题，而不是练出好模型。

判断"做对了"看三处：

1. 终端打印出一条 wandb 链接（匿名 run，7 天后自动删除），点开能看到损失曲线；
2. 生成checkpoints/checkpoint_epoch1.pth文件；
3. 日志里出现Validation Dice score: 0.xxxx这一行，说明验证集评估链路是通的。

三样都齐了，把--epochs 1换成--epochs 50重跑，正式训练就开始了。

补齐地基：环境与数据一次到位

环境：三个硬性条件

• Python ≥ 3.6，PyTorch ≥ 1.13（带 CUDA 版本）
• requirements.txt里其实只有 5 个库：matplotlib、numpy、Pillow、tqdm、wandb，装起来不慢
• 显存：6GB（如 GTX 1060）能跑通训练，24GB 可以放开 batch size 和分辨率；纯推理 CPU 即可

数据：两条铁律

如果你不用 Carvana 而是自己的数据，只需满足两点：

• 图像放data/imgs/，掩码放data/masks/，同名配对（Carvana 的掩码自带_mask后缀，utils/data_loading.py 里的CarvanaDataset已经处理了）；
• 两个目录都不要建子文件夹，数据加载器会贪婪地扫平整个目录。

掩码用单通道表示，不同像素值对应不同类别。加载时 BasicDataset 自动完成三件事：按--scale缩放、图像除以 255 归一化、掩码像素值映射成类别索引。你不需要写任何预处理代码。

调参：只动该动的

值得调的参数

别动的部分

优化器（RMSprop）、学习率调度器（ReduceLROnPlateau，以验证 Dice 为信号自动降速）、BCE + Dice 组合损失、梯度裁剪，这些都已在 train.py 的train_model里调好，新手改动收益为负。另外 train.py 里还藏了个保险：一旦OutOfMemoryError，自动开启梯度检查点后重跑，速度变慢但能继续。

每个 epoch 结束都会落盘到checkpoints/checkpoint_epoch{N}.pth，中途断电也不心疼。

检验与上线：从评估到部署

评估：只看 Dice 就够

分割任务的核心指标是 Dice 系数，越接近 1 越好，实现就在 utils/dice_score.py。训练时每隔若干 step 自动跑一轮验证并打印分数；项目作者在 Carvana 测试集（10 万+ 张图）上报告的 Dice 约为 0.988，可以当参照上限。

可视化：wandb 帮你盯着

训练启动后打开 wandb 链接，能看到损失曲线、验证 Dice、权重/梯度直方图，以及真值掩码与预测掩码的逐图对比。想长期留档的话设置环境变量WANDB_API_KEY关联自己的账号，否则匿名 run 七天后即删。

预测：一条命令出掩码

# 保存结果 python predict.py -i data/imgs/xxx.jpg -o mask_out.jpg # 多张图、只弹窗口看、不落盘 python predict.py -i a.jpg b.jpg --viz --no-save

常用附加参数：--model指定 checkpoint（默认找MODEL.pth）、--mask-threshold控制二值化阈值（默认 0.5）、--scale要和训练时一致。多类输出走argmax，单类走 sigmoid 过阈值，逻辑在 predict.py 里一眼可查。

赶时间不想训练的话，仓库通过 hubconf.py 提供了官方预训练权重的加载方式：

import torch net = torch.hub.load('milesial/Pytorch-UNet', 'unet_carvana', pretrained=True, scale=0.5)

scale只支持 0.5 和 1.0，对应两种分辨率的预训练版本。

部署：Docker 一条命令

仓库自带 Dockerfile，也可直接用官方镜像：

docker run -it --rm --shm-size=8g --ulimit memlock=-1 --gpus all milesial/unet

进容器后数据下载、训练、预测的用法和本地完全一致，适合给没有 Python 环境的服务器做推理服务。

踩坑速查

进阶方向

跑通之后，按这个顺序往下走比较踏实：

1. 换自己的数据完整走一遍：准备data/imgs+data/masks，按通道改n_channels（灰度图是 1），体会多类场景下--classes和mask_values的配合；
2. 对比上采样方式：--bilinear开与关各训 5 个 epoch，用验证 Dice 和掩码边缘的锐利度做对比，理解反卷积上采样的代价；
3. 导出 ONNX 部署：几行代码即可把训好的 checkpoint 导给 C++ / 移动端用：

import torch from unet import UNet net = UNet(n_channels=3, n_classes=2) sd = torch.load('checkpoints/checkpoint_epoch50.pth', map_location='cpu') sd.pop('mask_values', None) net.load_state_dict(sd) net.eval() torch.onnx.export(net, torch.randn(1, 3, 256, 256), 'unet.onnx', opset_version=11)

把上面三步做完，这套代码对你来说就不再是一个 demo，而是一个可以随时换数据、换类别复用的分割基座。

信息块

• 项目地址：gh_mirrors/py/Pytorch-UNet（gitcode 镜像）
• 模型出处：U-Net: Convolutional Networks for Biomedical Image Segmentation（Ronneberger 等，arXiv: 1505.04597）
• 关键代码：网络结构 · 训练入口 · 数据加载 · Dice 实现

【免费下载链接】Pytorch-UNetPyTorch implementation of the U-Net for image semantic segmentation with high quality images项目地址: https://gitcode.com/gh_mirrors/py/Pytorch-UNet