此处用的FasterRCNN模型使用的是B导的源码,读者可以去B站搜B导的视频进行了解和学习,视频中B导非常细心讲解了如何训练自己的数据集以及预测。
文章目录
- 前言
- 一、准备数据集
- 二、环境配置
- 2.1 基础环境
- 2.2 其他依赖包安装
- 2.3 预训练权重下载
- 二、训练数据集
- 三、预测
- 结束语
前言
此实验的整个流程参考了B导的博客:睿智的目标检测27——Pytorch搭建Faster R-CNN目标检测平台 以及另一个博主的:对比实验系列:Faster R-CNN配置及训练自己数据集
B导的源码:https://github.com/bubbliiiing/faster-rcnn-pytorch B站视频链接:https://www.bilibili.com/video/BV1BK41157Vs/
一、准备数据集
此源码数据集使用的是voc格式,因为我已经拥有yolo格式的数据集,故采用的是yolo格式转为voc格式。yolo2voc参考的这篇文章:将Yolo格式标注文件转换为VOC格式,如果是其他格式的可以自行查找数据集格式的相互转换的方式。
转换后得到数据集的如下分布:
ps:由于我已经对数据集进行划分了训练集和验证集,故直接生成的train.txt和val.txt。
然后使用源码中的voc_annotation.py生成2007_train.txt 2007_val.txt(如果你的情况和我一样是已经划分好训练集 验证集的,只需要生成2007_train.txt 2007_val.txt,就在voc_annotation.py中设置mode=2,类别设置为自己的类别。
classes_path:按着顺序将自己数据集的类别写入
如果你并未划分,想随机分可以直接将mode=0,直接随机划分并生成相应的2007_train.txt 2007_val.txt。这两个文件包含图片的位置信息、图片中包含的目标信息。
就这样,数据集准备完成。
二、环境配置
2.1 基础环境
建议新建一个环境
conda create -n rcnn python=3.8 -y
conda activate rcnn
pip install torch==1.9.1+cu111 torchvision==0.10.1+cu111 -f https://download.pytorch.org/whl/torch_stable.html
2.2 其他依赖包安装
进入项目源码,点击终端运行
pip install -r requirements.txt
注意:安装requirements.txt的时候会报一些错,经试验,我将这两个修改为≥即可
2.3 预训练权重下载
将预训练权重下载至项目里面,可以在github上进行下载。
二、训练数据集
进入train.py修改超参数,主要修改这些,其他超参数可以根据情况自行调整。
classes_path为自己数据集的标签路径
model_path为预训练权重路径
调整冻结训练和非冻结训练的batch、线程数num_workers
train_annotation_path和val_annotation_path为第一部分生成的图片索引文件
开始训练
python train.py
训练完成以后可以得到以下训练文件
三、预测
首先到frcnn.py中去修改class FRCNN类:
然后就是修改predict.py中的超参数进行预测
由于我是对整个文件夹进行预测,故mode选择
然后修改文件夹路径以及输出的路径
如果想要保存txt文件可以参考注释
结束语
浅浅记录训练两阶段模型FasterRCNN模型的过程