U-Net: 使用PyTorch进行语义分割
这是一个为Kaggle的Carvana图像遮罩挑战赛中的高清图像,使用PyTorch实现的定制化U-Net。
快速开始
不使用Docker
-
安装依赖项
pip install -r requirements.txt
- 下载数据并运行训练:
bash scripts/download_data.sh
python train.py --amp
使用Docker
curl https://get.docker.com | sh && sudo systemctl --now enable docker
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
sudo docker run --rm --shm-size=8g --ulimit memlock=-1 --gpus all -it milesial/unet
- 下载数据并运行训练:
bash scripts/download_data.sh
python train.py --amp
描述
该模型从头开始训练,使用了5k张图像,并在超过100k张测试图像上获得了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
usage: 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
标志启用。混合精度可以使模型使用更少的内存,并在使用FP16运算的现代GPU上更快。建议启用AMP。
预测
在训练模型并将其保存到MODEL.pth
后,你可以通过CLI轻松测试图像上的输出遮罩。
要预测单个图像并保存它:
python predict.py -i image.jpg -o output.jpg
要预测多个图像并不保存它们而是直接展示:
python predict.py -i image1.jpg image2.jpg --viz --no-save
> python predict.py -h
usage: 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 Ronne