U-Net：基于 PyTorch 的语义分割

针对 Kaggle 的 Carvana 图像掩码挑战赛，基于高分辨率图像，在 PyTorch 中对 U-Net 进行了定制化实现。

• 快速入门
  • 不使用 Docker
  • 使用 Docker
• 描述
• 用法
  • Docker
  • 训练
  • 预测
• Weights & Biases
• 预训练模型
• 数据

快速入门

不使用 Docker

1. 安装 CUDA

2. 安装 PyTorch 1.13 或更高版本

3. 安装依赖项

pip install -r requirements.txt

4. 下载数据并运行训练：

bash scripts/download_data.sh
python train.py --amp

使用 Docker

1. 安装 Docker 19.03 或更高版本：

curl https://get.docker.com | sh && sudo systemctl --now enable docker

2. 安装 NVIDIA 容器工具包：

distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
   && curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - \
   && curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update
sudo apt-get install -y nvidia-docker2
sudo systemctl restart docker

3. 下载并运行镜像：

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

4. 下载数据并运行训练：

bash scripts/download_data.sh
python train.py --amp

描述

该模型从 5000 张图像中从头开始训练，在超过 10 万张测试图像上获得了 0.988423 的 Dice 系数。

它可以轻松用于多类分割、人像分割、医学分割等场景。

用法

注意：请使用 Python 3.6 或更高版本

Docker

包含代码和依赖项的 Docker 镜像已在 DockerHub 上发布。 您可以使用 (docker >=19.03) 下载并进入容器：

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

驯练

> python train.py -h
用法：train.py [-h] [--epochs E] [--batch-size B] [--learning-rate LR]
                [--load LOAD] [--scale SCALE] [--validation VAL] [--amp]

在图像和目标掩码上训练 UNet

可选参数：
  -h, --help            显示此帮助消息并退出
  --epochs E, -e E      训练轮数
  --batch-size B, -b B  批量大小
  --learning-rate LR, -l LR
                        学习率
  --load LOAD, -f LOAD  从 .pth 文件加载模型
  --scale SCALE, -s SCALE
                        图像的下采样因子
  --validation VAL, -v VAL
                        用作验证集的数据百分比 (0-100)
  --amp                 使用混合精度

默认情况下，scale 为 0.5，因此如果您希望获得更好的结果（但会占用更多内存），可以将其设置为 1。

还可以通过 --amp 标志启用自动混合精度。混合精度 允许模型在最近的 GPU 上使用更少的内存并提高速度，因为它采用 FP16 算法。建议启用 AMP。

预测

在训练完模型并将其保存到 MODEL.pth 后，您可以通过命令行轻松地在自己的图像上测试输出掩码。

要预测单张图像并保存： python predict.py -i image.jpg -o output.jpg

要预测多张图像并在不保存的情况下显示它们： python predict.py -i image1.jpg image2.jpg --viz --no-save

> python predict.py -h
用法：predict.py [-h] [--model FILE] --input INPUT [INPUT ...] 
                  [--output INPUT [INPUT ...]] [--viz] [--no-save]
                  [--mask-threshold MASK_THRESHOLD] [--scale SCALE]

根据输入图像预测掩码

可选参数：
  -h, --help            显示此帮助消息并退出
  --model FILE, -m FILE
                        指定存储模型的文件
  --input INPUT [INPUT ...], -i INPUT [INPUT ...]
                        输入图像的文件名
  --output INPUT [INPUT ...], -o INPUT [INPUT ...]
                        输出图像的文件名
  --viz, -v             在处理图像时进行可视化
  --no-save, -n         不保存输出掩码
  --mask-threshold MASK_THRESHOLD, -t MASK_THRESHOLD
                        将掩码像素视为白色所需的最小概率值
  --scale SCALE, -s SCALE
                        输入图像的比例因子

您可以使用 --model MODEL.pth 指定要使用的模型文件。

Weights & Biases

训练进度可以使用 Weights & Biases 实时可视化。损失曲线、验证曲线、权重和梯度直方图以及预测的掩码都会被记录到该平台上。

启动训练时，控制台会打印一个链接。点击该链接即可进入您的仪表板。如果您已有 W&B 账户，可以通过设置 WANDB_API_KEY 环境变量将其关联起来。如果没有，则会创建一个匿名运行，该运行将在 7 天后自动删除。

预训练模型

适用于 Carvana 数据集的 预训练模型 已经提供。它也可以从 torch.hub 加载：

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

可用的比例因子为 0.5 和 1.0。

数据

Carvana 数据集可在 Kaggle 网站 上获取。

你也可以使用辅助脚本下载数据：

bash scripts/download_data.sh

输入图像和目标掩码应分别放置在 data/imgs 和 data/masks 文件夹中（请注意，由于数据加载器的特性，imgs 和 masks 文件夹内不应包含任何子文件夹或其他文件）。对于 Carvana 数据集，图像为 RGB 格式，掩码为黑白图像。

你可以使用自己的数据集，只要确保它能在 utils/data_loading.py 中正确加载即可。

---

Olaf Ronneberger、Philipp Fischer 和 Thomas Brox 的原始论文：

U-Net：用于生物医学图像分割的卷积网络

PyTorch-UNet 快速上手指南

PyTorch-UNet 是一个基于 PyTorch 实现的 U-Net 语义分割模型，最初为 Kaggle Carvana 图像掩膜挑战赛设计，适用于医学图像分割、人像分割等多类场景。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐) 或 macOS/Windows (需配置 CUDA 环境)
• Python: v3.6 或更高版本
• GPU: 支持 CUDA 的 NVIDIA 显卡 (可选，但强烈推荐用于训练)
• CUDA: 已安装并配置好对应的 CUDA 驱动
• PyTorch: v1.13 或更高版本

安装步骤

您可以选择直接使用 Docker 容器（推荐，环境隔离性好）或在本地环境中安装。

方案一：使用 Docker (推荐)

如果您已安装 Docker 和 NVIDIA Container Toolkit，这是最快捷的方式。

1. 拉取并运行镜像： 以下命令将下载镜像并启动一个包含所有依赖的容器：

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

2. 进入容器后，继续执行下方的“基本使用”步骤。

方案二：本地安装 (Without Docker)

如果您希望在本地环境中运行：

1. 安装 PyTorch： 访问 PyTorch 官网 获取适合您环境的安装命令。 国内用户推荐使用清华源加速安装：

   pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
   # 或者使用清华镜像
   pip install torch torchvision torchaudio --index-url https://pypi.tuna.tsinghua.edu.cn/simple
   

2. 安装项目依赖： 克隆本项目代码库后，进入目录并安装依赖：

   git clone https://github.com/milesial/PyTorch-UNet.git
   cd PyTorch-UNet
   pip install -r requirements.txt
   

基本使用

1. 数据准备

本项目默认使用 Carvana 数据集。您可以使用内置脚本自动下载：

bash scripts/download_data.sh

注意：下载后的图片应位于 data/imgs，掩膜位于 data/masks。若使用自定义数据集，请确保目录结构一致，并检查 utils/data_loading.py 中的加载逻辑。

2. 训练模型

使用以下命令开始训练。建议开启 --amp 参数以启用混合精度训练，可显著减少显存占用并提升速度。

python train.py --amp

常用参数说明：

• --epochs E: 训练轮数
• --batch-size B: 批次大小
• --scale SCALE: 图像缩放因子 (默认 0.5，设为 1 可获得更好效果但消耗更多显存)
• --load LOAD: 从 .pth 文件加载预训练权重继续训练

3. 预测与推理

训练完成后（或使用预训练模型），可对单张或多张图片进行分割预测。

加载官方预训练模型进行预测 (Python 代码示例)：

import torch
from unet import UNet

# 加载预训练模型 (scale 可选 0.5 或 1.0)
net = torch.hub.load('milesial/Pytorch-UNet', 'unet_carvana', pretrained=True, scale=0.5)
net.eval()

使用命令行进行预测：

• 预测单张图片并保存结果：

  python predict.py -i image.jpg -o output.jpg
  

• 预测多张图片并实时预览（不保存）：

  python predict.py -i image1.jpg image2.jpg --viz --no-save
  

• 指定特定模型文件进行预测：

  python predict.py -m MODEL.pth -i image.jpg -o output.jpg
  

预测参数说明：

• --model: 指定模型权重文件路径
• --mask-threshold: 掩膜阈值 (0-1)，决定像素被判定为前景的最小概率
• --viz: 可视化处理过程