【记录】使用Pytorch-UNet训练自己的数据集

今天回去重新调试半个月前调试的算法，发现很多地方已经忘了，决定写一篇博客把调试代码的细节部分记录下。

第一步:制作数据集

只需要将相同数量的图片和其掩码放置在data路径下的imgs和masks文件夹里，图片名需要和掩码标签相同。

第二步:修改 utils/dataset.py

需要将源代码下篮框标注的地方改成如图所示

第三步:修改 train.py

train.py 文件里有如图传参的地方，从上到下分别意思是：

• epoch，训练批次
• batch size，一次输入图片的数量
• learning rate，学习率
• load，加载预训练模型的路径
• scale，该值为0-1之间，设置越高，精度越高，但是资源占用越多
• validation，0-100之间，为验证集占数据中的百分比值