unet_pytorch图像分割算法模型

UNET

论文

U-Net: Convolutional Networks for Biomedical Image Segmentation

• https://arxiv.org/abs/1505.04597

模型结构

UNet（全名 U-Net）是一种用于图像分割的卷积神经网络（CNN）架构，UNet 的结构具有 U 形状，因此得名。

算法原理

U-Net 的核心原理如下：

1. 编码器（Contracting Path）：U-Net 的编码器由卷积层和池化层组成，用于捕捉图像的特征信息并逐渐减小分辨率。这一部分的任务是将输入图像缩小到一个低分辨率的特征图，同时保留有关图像内容的关键特征。
2. 中间层（Bottleneck）：在编码器和解码器之间，U-Net 包括一个中间层，通常由卷积层组成，用于进一步提取特征信息。
3. 解码器（Expansive Path）：U-Net 的解码器包括上采样层和卷积层，用于将特征图恢复到原始输入图像的分辨率。解码器的任务是将高级特征与低级特征相结合，以便生成分割结果。这一部分的结构与编码器相对称。

   

环境配置

Docker（方法一）

docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.1.0-centos7.6-dtk24.04-py310

docker run -it --network=host --name=unet --privileged -v /opt/hyhal:/opt/hyhal:ro --device=/dev/kfd --device=/dev/dri --ipc=host --shm-size=32G  --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -u root --ulimit stack=-1:-1 --ulimit memlock=-1:-1  image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.1.0-centos7.6-dtk24.04-py310

cd /path/unet-pytorch
pip install -r requirements.txt
pip install NumPy==1.23.5

Dockerfile（方法二）

docker build --no-cache -t unet:latest .
docker run -dit --network=host --name=unet --privileged -v /opt/hyhal:/opt/hyhal:ro --device=/dev/kfd --device=/dev/dri --ipc=host --shm-size=16G --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -u root --ulimit stack=-1:-1 --ulimit memlock=-1:-1 unet:latest
docker exec -it unet /bin/bash

cd /path/unet-pytorch
pip install -r requirements.txt
pip install NumPy==1.23.5

Anaconda（方法三）

关于本项目DCU显卡所需的特殊深度学习库可从光合开发者社区下载安装。

DTK驱动：dtk24.04
python：python3.10
apex：1.1.0+0dd7f68.abi0.dtk2404.torch2.1
torch：2.1.0+git00661e0.abi0.dtk2404
torchvision：0.16.0+gitc9e7141.abi0.dtk2404.torch2.1

Tips：以上dtk驱动、python等DCU相关工具版本需要严格一一对应

其它非深度学习库参照requirements.txt安装：

pip install -r requirements.txt
pip install NumPy==1.23.5

数据集

Carvana

• AIDatasets / project-dependency / Carvana · 极狐GitLab

下载命令：

bash scripts/download_data.sh

项目中已提供用于试验训练的迷你数据集，训练数据目录结构如下，用于正常训练的完整数据集请按此目录结构进行制备：

│ ── data
    │   ├── imgs
    │   ├────── fff9b3a5373f_15.jpg
    │   ├────── fff9b3a5373f_16.jpg
    │   └── masks
    │   ├────── fff9b3a5373f_15.gif
    │   ├────── fff9b3a5373f_16.gif

训练

关闭wandb

wandb disabled

单机单卡

python train.py

单机多卡

python -m torch.distributed.launch --nproc_per_node 4 train_ddp.py

推理

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

result

精度

测试数据：test data，使用的加速卡:Z100L。（采用iou系数）

根据测试结果情况填写表格：

Unet	精度
Carvana	0.976

应用场景

算法类别

图像分割

热点应用行业

医疗,交通,家居

预训练权重

SCNet AIModels - milesial-pytorch-unet-3.0

源码仓库及问题反馈

• ModelZoo / unet_pytorch · GitLab

参考资料

• 推理https://github.com/milesial/Pytorch-Unet