U-Net训练自己的数据集(pytorch)

1.下载U-Net训练网络

网络：链接：https://pan.baidu.com/s/1c-Fumks0bsYgo4BOzTkOzQ 
提取码：njwa

预训练权重：链接：https://pan.baidu.com/s/1_NCV5Slr9uPKnFn9eX4AUA 
提取码：xz3y

 

2.数据集格式

选用voc数据格式

 

 

 

3.数据集设置

（1）将自己的图片和标签文件放到VOCdevkit下相应文件夹内

（2）运行voc_annotation.py自动生成训练集和测试集.txt

 

 

4.预训练权重

（1）将下载好的预训练权重放到model_data

 （2）修改train.py中权重文件路径

4.开始训练

根据注释内容修改对应的参数

num_classes指向检测类别的个数+1

 

（2）训练后的权重文件和日志文件位于logs中

5.训练结果预测

 （1）训练结果预测需要用到两个文件，分别是unet.py和predict.py。
我们首先需要去unet.py里面修改model_path以及num_classes，这两个参数必须要修改。

model_path指向训练好的权值文件，在logs文件夹里。
num_classes指向检测类别的个数+1。

修改unet.py中相应的文件路径，和类别数

（2）运行prediect.py文件

参考：Predict-（训练必看）利用训练好的模型进行预测_哔哩哔哩_bilibili