PyTorch使用U-Net进行图像语义分割训练和测试代码

# U-Net：使用PyTorch进行语义分割  在PyTorch中对[U-Net](https://arxiv.org/abs/1505.04597)的自定义实现，用于解决Kaggle的[Carvana图像遮挡挑战](https://www.kaggle.com/c/carvana-image-masking-challenge)，该挑战涉及对高清图像进行语义分割。 - [快速开始](#quick-start) - [无需Docker](#without-docker) - [使用Docker](#with-docker) - [描述](#description) - [用法](#usage) - [Docker](#docker) - [训练](#training) - [预测](#prediction) - [Weights & Biases](#weights--biases) - [预训练模型](#pretrained-model) - [数据](#data) ## 快速开始 1. [安装CUDA](https://developer.nvidia.com/cuda-downloads) 2. [安装PyTorch 1.13或更高版本](https://pytorch.org/get-started/locally/) 3. 安装依赖项 ```bash pip install -r requirements.txt ``` 4. 下载数据并运行训练: ```bash bash scripts/download_data.sh python train.py --amp ``` ## 描述 该模型是从头开始使用5k张图像进行训练的，对100k多张测试图像获得了0.988423的[Dice系数](https://en.wikipedia.org/wiki/S%C3%B8rensen%E2%80%93Dice_coefficient)。 该模型可以轻松用于多类别分割、人像分割、医学分割等等... ## 用法 **注意: 使用Python 3.6或更新版本** ### 训练 ```console > python train.py -h usage: train.py [-h] [--epochs E] [--batch-size B] [--learning-rate LR] [--load LOAD] [--scale SCALE] [--validation VAL] [--amp] Train the UNet on images and target masks optional arguments: -h, --help show this help message and exit --epochs E, -e E Number of epochs --batch-size B, -b B Batch size --learning-rate LR, -l LR Learning rate --load LOAD, -f LOAD Load model from a .pth file --scale SCALE, -s SCALE Downscaling factor of the images --validation VAL, -v VAL Percent of the data that is used as validation (0-100) --amp Use mixed precision ``` 默认情况下，`scale`为0.5，如果希望获得更好的结果（但使用更多内存），请将其设置为1。 还可以通过`--amp`标志启用自动混合精度。[混合精度](https://arxiv.org/abs/1710.03740)允许模型使用更少的内存，并在最新的GPU上更快地执行，因为它使用FP16算法。建议启用AMP。 ### 预测 在训练模型并保存为`MODEL.pth`之后，可以通过CLI轻松测试图像的输出掩模。 要预测单个图像并保存它: `python predict.py -i image.jpg -o output.jpg` 要预测多个图像并在不保存它们的情况下显示它们: `python predict.py -i image1.jpg image2.jpg --viz --no-save` ```console > 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] Predict masks from input images optional arguments: -h, --help show this help message and exit --model FILE, -m FILE Specify the file in which the model is stored --input INPUT [INPUT ...], -i INPUT [INPUT ...] Filenames of input images --output INPUT [INPUT ...], -o INPUT [INPUT ...] Filenames of output images --viz, -v Visualize the images as they are processed --no-save, -n Do not save the output masks --mask-threshold MASK_THRESHOLD, -t MASK_THRESHOLD Minimum probability value to consider a mask pixel white --scale SCALE, -s SCALE Scale factor for the input images ``` 可以通过`--model MODEL.pth`指定要使用的模型文件。 ## Weights & Biases 训练进度可以通过[Weights & Biases](https://wandb.ai/)实时可视化。 损失曲线、验证曲线、权重和 梯度直方图以及预测的掩码都将记录到该平台。 在启动训练时，将在控制台中打印一个链接。 单击它即可转到您的仪表板。 如果您已经有W&B帐户，可以通过设置`WANDB_API_KEY`环境变量来链接它。 如果没有，它将创建一个匿名运行，将在7天后自动删除。 ## 预训练模型 可以在[此处](https://github.com/milesial/Pytorch-UNet/releases/tag/v3.0)获取Carvana数据集的[预训练模型](https://github.com/milesial/Pytorch-UNet/releases/tag/v3.0)。 还可以从torch.hub加载它: ```python net = torch.hub.load('milesial/Pytorch-UNet', 'unet_carvana', pretrained=True, scale=0.5) ``` 可用的比例为0.5和1.0。 ## 数据 Carvana数据可在[Kaggle网站](https://www.kaggle.com/c/carvana-image-masking-challenge/data)上找到。 还可以使用辅助脚本下载数据: ``` bash scripts/download_data.sh ``` 输入图像和目标掩码应分别位于`data/imgs`和`data/masks`文件夹中（请注意，`imgs`和`masks`文件夹不应包含任何子文件夹或其他文件，由于贪婪的数据加载器）。 对于Carvana，图像是RGB的，掩码是黑白的。 只要确保在`utils/data_loading.py`中正确加载数据，就可以使用自己的数据集。 --- Original paper by Olaf Ronneberger, Philipp Fischer, Thomas Brox: [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://arxiv.org/abs/1505.04597)