训练自己的物体检测模型

前言

在上一篇博客中，介绍了怎样利用SSD实现物体检测，我们使用的是官方的模型。在本文中，将继续介绍在此基础上如何实现迁移学习，训练自己的数据集。

1. 准备数据集

1. 数据集来源一般有两个：使用labelImg自己标注，或者使用网上开源的数据集。本文使用的是开源的VOC 2007，总共有20个分类。数据集包括5001张图片，以及对应的.xml文件，分别存放在image和annotation文件夹中。
2. 本文通过.csv文件生成.record文件，如果习惯直接使用.xml文件生成.record文件，可以跳过这一步。转换程序xml_to_csv.py已准备好，直接运行即可在当前路径下生成labels.csv，顺便说一下，本项目的.py程序都配备了一个同名的.bat启动程序，直接双击后者也可以，转换效果如图：
   
   注意： 在准备数据阶段，可以同步准备测试集。在xml_to_csv.py中需要修改如下两个参数，对应测试集的.xml路径和输出的.csv文件名：
   
3. 生成.record文件：
   程序也已准备好：generate_tfrecord.py和generate_tfrecord.bat，我们只需要指定.bat文件中的三个参数即可运行。分别是：
   csv_input: 标签文件csv
   images_dir: 图片存放路径
   output_path: 输出的文件
   
   注意： 在generate_tfrecord.py有一段代码将20个类别转换为数字，如果改变数据集，这里是需要做相应调整的：
   
4. 准备.pbtxt文件
   这个文件用来解码数字和类别，要和之前generate_tfrecord.py中的对应关系一致：
   

2. 准备预训练好的模型

1. 下载模型
   tensorflow官方提供了很多训练好的模型，我们可以选择一个下载：
   
   下载后解压，.pb文件是在之前的博客中用到过的，现在我们要继续训练，需要的是.ckpt文件，将他们放到train目录下
   
2. 设置配置文件
   在./object_detection/samples/configs/下有很多配置文件，对应各个模型，此处我们选择ssd_mobilenet_v1_coco.config，将其复制到当前路径，然后修改几个参数：

   num_classes: 20
   batch_size: 16 # 计算力强的话可以增大
   num_steps: 2500 # 训练步数，达到就停止，即使再次运行
   fine_tune_checkpoint: “./training/model.ckpt”

3. 运行train.bat，第一个参数固定，第二个是配置文件，第三个是模型输出路径
   
4. 运行结果
   
   
   此时，我们还可以改大配置文件中的num_steps，继续训练

3. 输出模型

我们之前输出的是.ckpt模型，为了方便预测，我们可以先把模型转化为.pb格式，运行generate_tfrecord.bat，第一个参数固定，第二个是配置文件，第三个是选择要转化的模型，第四个是新模型的输出路径

4. 预测

预测的程序和之前的一样，只删掉了现在模型的部分，直接用我们自己的模型