UNet - 实现模型的训练与预测 - 入门1

最新推荐文章于 2024-04-25 12:29:32 发布

正在努力的鼠鼠

最新推荐文章于 2024-04-25 12:29:32 发布

阅读量1.4k

点赞数 1

文章标签：服务器机器学习

本文链接：https://blog.csdn.net/qq_52589927/article/details/131621019

版权

前言

本篇文章，是讲如何在服务器上跑 unet 模型，正文中会放许多跳转链接，都是有一定参考价值的

UNets是什么

UNet是一种用于图像分割的深度学习网络模型。

UNet的结构类似于一个U形的结构，因此得名UNet。编码器部分负责提取图像的高级特征表示，而解码器部分则通过上采样和跳跃连接将特征逐渐恢复到原始图像尺寸，并生成像素级别的分割结果。

UNet的结构使得它能够在处理医学图像等任务中取得出色的性能。通过跳跃连接，UNet能够保留更多的细节和空间信息，并且在训练过程中，可以更好地学习到不同尺度的特征表示。这使得UNet在处理边缘、小物体和图像内部结构等方面具有优势。

由于其简单有效的结构和出色的性能，在图像分割任务中，特别是医学图像分割领域，UNet已经成为一种广泛使用的模型。

（这段看不懂没关系，可以在网上找找详细介绍）

简单实现

在熟悉了UNet之后，第一步，先想办法把 UNet 跑起来看看

进入服务器

有关如何使用服务器跑代码，请看这篇文章

服务器使用指南https://blog.csdn.net/qq_52589927/article/details/131608338?spm=1001.2014.3001.5501

准备数据集

GitHub里面的数据集，我发现不能用，所有我又重新找了。在 AI studio 中，链接放在这里，自行下载（300MB），不是很大，初学者用正好

我是数据集https://aistudio.baidu.com/aistudio/datasetdetail/184683然后需要将下载好的数据上传到服务器中，流程在文章开头的 “服务器使用指南” 中会讲，

我用记事本打开了教程中下载数据的脚本文件，在里面我知道了，数据要分别存在 img 和 masks 文价夹中

所以，要把数据用 fz 分别拖进各自的文件夹，数据集这就准备好了

找一个教程

https://github.com/milesial/Pytorch-UNethttps://github.com/milesial/Pytorch-UNet把上面的代码下载下来（如图），解压缩后传到服务器（怎么上传服务器后文链接里有）

接下来看一看 README.md - Quick start - Without Docker

租的服务器一般自带 CUDA 和 Pytorch，所以直接开始安装依赖，终端输入

pip install -r requirements.txt

这个 requirement.txt 中有matplotlib==3.6.2、numpy==1.23.5、Pillow==9.3.0、tqdm==4.64.1、wandb==0.13.5

1. 训练

教程中，接下来需要运行一个脚本，里面可以下载数据集，但是咱们不用那个数据，所以直接开始训练，cd到有 train.py 的文件夹，终端输入

python train.py --amp

效果如图所示

下方是Linux命令教学

用其中的命令可以对模型的一些参数进行调整

> python train.py -h
用法:

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

可选项:
-h, --help 打开帮助文档
--epochs E, -e E epochs数量
--batch-size B, -b B Batch size
--learning-rate LR, -l LR 学习率
--load LOAD, -f LOAD 从 .pth 文件加载模型
--scale SCALE, -s SCALE 图像的缩小因子
--validation VAL, -v VAL 用作验证的数据百分比（0-100）
--amp 使用混合精度

2. 预测

输入一张图片，用训练好的模型，预测这张图片，终端输入

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

随便找一张图，上传到服务器，例 101.jpg

还需要准备好刚刚训练的模型权重，通常在checkpoints文件夹中，接下来可以看到，里面有5个 .pth 文件，选择最后一个 epoch5，这是最后一轮训练的权重，通常来说是最好的效果。

-i 后面是输入数据的路径，-o 后面是输出文件，-m 后面是权重文件的路径，在终端输入（举例）

python predict.py -i /root/lzc_study/unet/data/101.jpg -o 101output.jpg -m /root/lzc_study/unet/checkpoints/checkpoint_epoch5.pth

> 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 输入图像的比例因子