【源码】基于Python和PyTorch的视网膜血管分割模型UNet

项目简介

本项目实现了基于Python和PyTorch的视网膜血管分割模型U-Net。U-Net是常用于图像分割任务的卷积神经网络架构，尤其适用于生物医学图像分割。本项目聚焦于视网膜血管分割，支持单GPU和多GPU训练，还提供简易预测脚本。

项目的主要特性和功能

1. U-Net模型实现：基于PyTorch实现U-Net模型，用于视网膜血管分割。
2. 多GPU训练支持：可使用多GPU进行分布式训练，提升训练效率。
3. 自定义数据集：通过my_dataset.py读取DRIVE数据集。
4. 训练和验证：提供train.py和train_multi_GPU.py进行单GPU和多GPU训练，predict.py用于模型预测。
5. 数据集统计：compute_mean_std.py统计数据集各通道均值和标准差，便于数据预处理。

安装使用步骤

环境配置

• 安装Python 3.6/3.7/3.8。
• 安装PyTorch 1.10。
• 运行pip install -r requirements.txt安装项目依赖。

数据准备

• 下载DRIVE数据集，放置在项目根目录下的data文件夹中。
• 运行compute_mean_std.py统计数据集的均值和标准差。

模型训练

• 单GPU训练：运行python train.py --data-path=./data。
• 多GPU训练：运行torchrun --nproc_per_node=8 train_multi_GPU.py --data-path=./data。
• 指定GPU设备：在命令前加上CUDA_VISIBLE_DEVICES=0,3，如CUDA_VISIBLE_DEVICES=0,3 torchrun --nproc_per_node=2 train_multi_GPU.py。

模型预测

使用训练好的权重进行预测，运行python predict.py --weights_path=./weights/model.pth --data-path=./data。

下载地址

点击下载 【提取码: 4003】【解压密码: www.makuang.net】