自己是计算机方向研一的学生,打算做的方向是图像分割方面,又因为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.总结
这方面知识还是很薄弱,还得多学习点这方面的知识,争取能够慢慢开始自己调整一些简单的参数,能够让模型得到一些更好的结果。