用Tensorflow的object detection api训练自己的数据

 之前的文章有用caffe去训练一个物体检测网络MobileNetV2-SSDLite，然后这里是用Tensorflow的object detection api训练MobileNetV2-SSDLite。

---

数据集准备

 数据集的准备可以直接参考目标检测Tensorflow object detection API之构建自己的模型或者自行搜索相应教程，最终应该有两个tfrecord文件（分别用于训练和测试）。

开发环境

 这里我的环境是python 3.7+Tensorflow 1.13.1。

object detection api 下载与安装

 使用git clone https://github.com/tensorflow/models.git下载所需的代码及模型文件。

 Windows下的话，先在https://github.com/google/protobuf/releases下载相应版本的protobuf，解压后将bin文件夹中的protoc.exe复制到models文件夹的research目录下，在research目录执行下面的代码，将object_detection/protos下的.proto文件转换成.py文件：

import os

for each in os.listdir('object_detection/protos'):
    if each.endswith('proto'):
        os.system('protoc object_detection/protos/%s --python_out=.' % each)

 Linux下的话可以参考linux下安装tensorflow object detection API以及 安装过程问题解决。

 接着在research和research/slim目录下分别执行python setup.py install以执行安装，然后再把research和research/slim目录添加到环境变量中去。

 为了验证安装成功与否，可以执行python object_detection/builders/model_builder_test.py，成功的话会输出OK。

训练自己的数据

 为了整理方便，创建一个文件夹目录如下：

├── MobileNetV2-SSDLite
  ├── data
  ├── models
  ├── training

 其中data用于存放训练及测试数据，models就是之前下载的代码及模型，training则存放模型文件及之后的训练过程产生的数据。

 将之前生成的两个tfrecord文件移到data文件夹下之后，再新建一个my_label_map.pbtxt文件，根据实际情况写入（这里因为我的数据只有一个类别所以就只有一个item）：

item {
    name : "tfeature"
    id : 1
}

 接着在object_detection文件夹中的samples/config路径下，找到配置文件ssdlite_mobilenet_v2_coco.config，复制到training文件夹中，并做如下修改：

• num_classes修改为1（当然根据你自己的数据作修改）
• batch_size根据电脑/服务器配置，可以适当调高或者调低
• 两个input_path分别设置成data/train.tfrecord和data/test.tfrecord
• 两个label_map_path均设置成data/my_label_map.pbtxt
• 从头训练的话，第158、159行的

    fine_tune_checkpoint: "PATH_TO_BE_CONFIGURED/model.ckpt"
  	fine_tune_checkpoint_type:  "detection"

  这2行需要删除或注释掉

 到这就可以开始训练了，但为了输出loss信息，先在model_main.py文件的import区域之后添加tf.logging.set_verbosity(tf.logging.INFO)，接着在MobileNetV2-SSDLite目录下执行以下命令：

python models/research/object_detection/model_main.py \
    --pipeline_config_path=training/ssdlite_mobilenet_v2_coco.config \
    --model_dir=training \
    --num_train_steps=400000 \
    --num_eval_steps=20 \
    --alsologtostderr

 接着就开始训练了，正常的话可以应该有如下的输出：

训练过程

 此外也可以用tensorboard --logdir=training可视化训练过程：

可视化训练过程

输出模型

 训练完成之后，可以通过下面的命令输出模型：

python models/research/object_detection/export_inference_graph.py \
    --input_type=image_tensor \
    --pipeline_config_path=./training/ssdlite_mobilenet_v2_coco.config \
    --trained_checkpoint_prefix=../Detection/backup/model.ckpt-39160 \
    --output_directory=inference_graph

 其中的路径根据自己的实际情况修改即可。执行完上面的命令之后会生成一个inference_graph文件夹，其下有saved_model文件夹和checkpoint、frozen_inference_graph.pb等文件，其中.pb文件将是接下来要用到的最重要的文件。

测试模型

 用jupyter打开object_detection目录下的object_detection_tutorial.ipynb文件，修改一些实际路径再运行就可以得到测试结果了。