基于pytorch的U-net模型简单复现

自己是计算机方向研一的学生，打算做的方向是图像分割方面，又因为u-net也是做医学图像方面很基础的网络，于是写个文章记录下自己复现别人代码的过程，学习借鉴一下。同时有错误也请各位大佬批评指正。

关于U-net具体的代码以及相关过程参考了以下这篇博客：
Pytorch入门——用UNet网络做图像分割

---

---

整体的文件目录

---

一、一些参数的调整

1、先把action中的默认值设置为train，随后跑完得到模型权重后，记得设置default为test，开始验证。
2、数据集里面的图片的尺寸较大，所以为了避免爆显存，这里把里面batch_size设置为了4。但GPU还是占用了不少的空间（自己的显卡是RTX3060,显存是12G的）。

二、训练过程

1.训练的结果

最后的损失：

跑完训练的文件目录如下：

保存了第20轮训练完成的模型pth(自己因为是复现别人的代码以及程序，就只跑了20个epoch)

2.遇到的问题

训练完成后跑test的时候遇到以下问题

这是因为没加入模型的路径。

解决如下：
在main.py的主函数里面的weight加入你训练好的模型权重的路径。

举例如下：

3.验证结果

其中val下 带有mask后缀的图片可以理解为标签图片。

---

这里也把test跑出来的结果和真实的mask标签图片做了个对比。

4.总结

这方面知识还是很薄弱，还得多学习点这方面的知识，争取能够慢慢开始自己调整一些简单的参数，能够让模型得到一些更好的结果。