前言
在上一篇博客中,介绍了怎样利用SSD实现物体检测,我们使用的是官方的模型。在本文中,将继续介绍在此基础上如何实现迁移学习,训练自己的数据集。
1. 准备数据集
- 数据集来源一般有两个:使用labelImg自己标注,或者使用网上开源的数据集。本文使用的是开源的VOC 2007,总共有20个分类。数据集包括5001张图片,以及对应的
.xml
文件,分别存放在image
和annotation
文件夹中。 - 本文通过
.csv
文件生成.record
文件,如果习惯直接使用.xml
文件生成.record
文件,可以跳过这一步。转换程序xml_to_csv.py
已准备好,直接运行即可在当前路径下生成labels.csv
,顺便说一下,本项目的.py
程序都配备了一个同名的.bat
启动程序,直接双击后者也可以,转换效果如图:
注意: 在准备数据阶段,可以同步准备测试集。在xml_to_csv.py
中需要修改如下两个参数,对应测试集的.xml
路径和输出的.csv
文件名:
- 生成
.record
文件:
程序也已准备好:generate_tfrecord.py
和generate_tfrecord.bat
,我们只需要指定.bat
文件中的三个参数即可运行。分别是:
csv_input:
标签文件csv
images_dir:
图片存放路径
output_path:
输出的文件
注意: 在generate_tfrecord.py
有一段代码将20个类别转换为数字,如果改变数据集,这里是需要做相应调整的:
- 准备
.pbtxt
文件
这个文件用来解码数字和类别,要和之前generate_tfrecord.py
中的对应关系一致:
2. 准备预训练好的模型
- 下载模型
tensorflow官方提供了很多训练好的模型,我们可以选择一个下载:
下载后解压,.pb
文件是在之前的博客中用到过的,现在我们要继续训练,需要的是.ckpt
文件,将他们放到train
目录下
- 设置配置文件
在./object_detection/samples/configs/
下有很多配置文件,对应各个模型,此处我们选择ssd_mobilenet_v1_coco.config
,将其复制到当前路径,然后修改几个参数:num_classes: 20
batch_size: 16 # 计算力强的话可以增大
num_steps: 2500 # 训练步数,达到就停止,即使再次运行
fine_tune_checkpoint: “./training/model.ckpt” - 运行
train.bat
,第一个参数固定,第二个是配置文件,第三个是模型输出路径
- 运行结果
此时,我们还可以改大配置文件中的num_steps
,继续训练
3. 输出模型
我们之前输出的是.ckpt
模型,为了方便预测,我们可以先把模型转化为.pb
格式,运行generate_tfrecord.bat
,第一个参数固定,第二个是配置文件,第三个是选择要转化的模型,第四个是新模型的输出路径
4. 预测
预测的程序和之前的一样,只删掉了现在模型的部分,直接用我们自己的模型