使用tensorflow_object_detection_api训练自定义模型

发表于 更新于 阅读次数：

近期研究目标对象检测和识别，发现谷歌开源的基于tensorflow的object detection api模型效果不错，于是git clone下来测试一下。下面记录我在debian linux上安装配置object detection api模型，构建自定义数据集，训练和测试object detection api模型的过程，以及整个过程中遇到的一下问题，需要注意的事项。

准备工作

docker安装

TensorFlow 程序在 GPU 上的运行速度通常要比在 CPU 上快得多，在系统满足NVIDIA 软件要求的前提下，推荐使用支持 GPU 的 TensorFlow。

官方推荐使用Docker简化TensorFlow的GPU支持配置，这样只需要linux主机安装好Nvidia GPU驱动即可。

  1. 安装docker
    安装docker的方法可参考链接1。

  2. 安装nvidia-docker

1
2
3
4
5
6
7
8
9
# Add the package repositories
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/debian9/nvidia-docker.list | \
  sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update

# Install nvidia-docker2 and reload the Docker daemon configuration
sudo apt-get install -y nvidia-docker2
sudo service docker restart
  1. 下载最新支持GPU和Python3的tensorflow映像[2]
1
docker pull tensorflow/tensorflow:latest-gpu-py3
  1. 测试tensorflow映像
1
2
docker run --runtime=nvidia -it --rm tensorflow/tensorflow:latest-gpu-py3 \
    python -c "import tensorflow as tf; print(tf.contrib.eager.num_gpus())"

下载TFModel库及其相关库

1
2
3
4
# 基于tensorflow的模型和例子
git clone https://github.com/tensorflow/models.git
# raccoon数据集，可参考其中生成TFRecord格式数据的方法
git clone https://github.com/datitran/raccoon_dataset.git

安装Python虚拟环境管理工具

1
2
3
4
# 具体安装配置过程参见链接3
sudo pip3 install virtualenv
sudo pip3 install virtualenvwrapper
mkvirtualenv object-detection --python=/usr/bin/python3# 创建python3虚拟环境

数据标注

以分类几何形状为例，收集相关图片后，需要对它们进行标注。推荐使用 LabelImg 进行标注，生成的文件是 PASCAL VOC 的 xml 格式。这个工具还可以加载标注文件，检查标注结果[3]。

1
2
3
4
5
6
7
8
9
git clone https://github.com/tzutalin/labelImg.git
workon object-detection
# 启动python3虚拟环境
workon object-detection
# 安装配置labelImg
sudo apt-get install pyqt5-dev-tools
sudo pip3 install -r requirements/requirements-linux-python3.txt
make qt5py3
python3 labelImg.py

除了标注图片，还需创建一个 .pbtxt 文件用来说明标注的分类。例如：

1
2
3
4
5
6
7
8
9
10
11
12
item {
  id: 1
  name: 'circle'
}
item {
  id: 2
  name: 'square'
}
item {
  id: 3
  name: 'triangle'
}

需要特别注意以下两点：

  • name 不支持直接写中文，需要 UTF-8 编码
  • id 从 1 开始编号，不能从 0 开始编号。

数据准备

标注好的数据是图片文件（.jpg）和标注文件（.xml），而 TensorFlow 不能直接用这些数据作为输入，还需要转成 TFRecord 格式。可采用两种方法生成TFRecord格式数据：

  • tensorflow/models中方法

项目中使用create_pascal_tf_record.py, create_pet_tf_record.py生成TFrecord 格式数据

  • raccoon_dataset中方法

项目中使用xml_to_csv.py把xml合并成一个CSV文件，使用split labels.ipynb 随机划分训练集和测试集，以及使用generate_tfrecord.py 生成相对应的 TFRecord 训练集和测试集。

1
2
3
4
5
6
7
8
9
workon object-detection
# 将生成geometry_labels.csv
python xml_to_csv.py
# 打开split labels.ipynb生产训练集train_labels.csv和测试集test_labels.csv
jupyter notebook
# 创建训练用TFRecord文件:
python generate_tfrecord.py --csv_input=data/train_labels.csv  --output_path=training/geometry_train.record
# 创建测试用TFRecord文件:
python generate_tfrecord.py --csv_input=data/test_labels.csv  --output_path=training/geometry_test.record

第二种方法比较直观，因此我采用第二种方法生产TFRecord数据。

安装配置

安装TFModel依赖

1
2
3
4
5
6
7
8
9
10
11
12
workon object-detection
# For CPU
pip install tensorflow
# For GPU
pip install tensorflow-gpu

pip install --user Cython
pip install --user contextlib2
pip install --user pillow
pip install --user lxml
pip install --user jupyter
pip install --user matplotlib

安装COCO API

1
2
3
4
git clone https://github.com/cocodataset/cocoapi.git
cd cocoapi/PythonAPI
make
cp -r pycocotools <path_to_tensorflow>/models/research/

编译Protobuf

1
2
3
4
5
6
# From tensorflow/models/research/
wget -O protobuf.zip https://github.com/google/protobuf/releases/download/v3.0.0/protoc-3.0.0-linux-x86_64.zip
unzip protobuf.zip

# From tensorflow/models/research/
./bin/protoc object_detection/protos/*.proto --python_out=.

添加库到PYTHONPATH

1
2
3
# From tensorflow/models/research/
# 每次调用TFModel的Object Detection API之前都要设置
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim

安装测试

1
2
# if all tests is ok, the installation is no problem.
python object_detection/builders/model_builder_test.py

模型训练

数据准备

  1. 将label_map文件、训练用TFRecord文件、测试用TFRecord文件复制到tensorflow/models/research/object_detection/data文件夹下。
1
2
3
4
5
# 推荐的数据目录结构
+ data
  - geometry.pbtxt
  - geometry_test.record
  - geometry_train.record
  1. 下载COCO预训练模型用于迁移学习
1
2
3
# From tensorflow/models/research/object-detection
wget http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_coco_2018_01_28.tar.gz
tar -xvf ssd_mobilenet_v1_coco_2018_01_28.tar.gz
  1. 修改解压后模型文件夹中pipeline.config

解压ssd_mobilenet_v1_coco_2018_01_28模型后会看到一个 .config 文件，里面包含有模型的参数，训练的参数，评估的参数等。这里需要修改到的有，

  • 模型参数中的 num_classes，改成你的类别数，
  • 训练参数中的 fine_tune_checkpoint，采用迁移学习，这里路径的指向刚才下载的Pre-train模型，比如 ssd_mobilenet_v1_coco_11_06_2017/model.ckpt
  • train_input_reader 下面的 input_path，改成你的训练数据，例如 data/train.record。label_map_path，改成你的 pbtxt 文件路径，例如 data/object.pbtxt
  • eval_input_reader 下面的 input_path，也需要改成你的测试集，例如 data/test.record。同样，label_map_path，也改成你的 pbtxt 文件路径，例如 data/object.pbtxt
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
//pipeline.config修改示例
model {
  ssd {
    num_classes: 3
    ...
  }
}
train_config {
  ...
  fine_tune_checkpoint: "object_detection/ssd_mobilenet_v1_coco_2018_01_28/model.ckpt"
  from_detection_checkpoint: true
  num_steps: 50000
}
train_input_reader {
  label_map_path: "object_detection/data/geometry.pbtxt"
  tf_record_input_reader {
    input_path: "object_detection/data/geometry_train.record"
  }
}
eval_input_reader {
  label_map_path: "object_detection/data/geometry.pbtxt"
  shuffle: false
  num_epochs: 1
  num_readers: 1
  tf_record_input_reader {
    input_path: "object_detection/data/geometry_test.record"
  }
  sample_1_of_n_examples: 1
}

训练数据

在GPU上训练TFModel，与在CPU上训练TFModel相比，要快五倍左右，因此推荐在GPU上训练TFModel。

  1. 启动docker
1
2
# 启动docker，并使用-v参数将目录/home/$USER挂载到容器上
docker run --runtime=nvidia -v /home/$USER:/home/$USER -it tensorflow/tensorflow:latest-gpu-py3 bash
  1. 执行训练
1
2
3
4
# From tensorflow/models/research/
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim
# 开始训练
python object_detection/model_main.py --pipeline_config_path=object_detection/ssd_mobilenet_v1_coco_2018_01_28/pipeline.config --model_dir=object_detection/ssd_mobilenet_v1_coco_2018_01_28/saved_model/ --num_train_steps=50000 --alsologtostderr
  1. 使用tensorboard查看训练进度
1
2
# From tensorflow/models/research/
tensorboard --logdir=object_detection/ssd_mobilenet_v1_coco_2018_01_28/saved_model/

在CPU上训练只需执行第二步和第三步。

导出模型

1
2
# From tensorflow/models/research/
python object_detection/export_inference_graph.py --input_type=image_tensor --pipeline_config_path=object_detection/ssd_mobilenet_v1_coco_2018_01_28/pipeline.config --trained_checkpoint_prefix=object_detection/ssd_mobilenet_v1_coco_2018_01_28/saved_model/model.ckpt-50000 --output_directory=../../../../ssd_mobilenet_v1_coco_2018_01_28

运行完命令后模型就导出到 ssd_mobilenet_v1_coco_2018_01_28 文件夹中，其中的frozen_inference_graph.pb即是所需模型。 需要注意的是，参数中的 –trained_checkpoint_prefix 是需要指定到单个模型的，例如 model.ckpt-50000，这个50000就是训练了 50000 步后自动保存模型。

参考文献

  1. https://huangwang.github.io/2018/10/18/Debian-Linux下安装Docker的方法/， by jack huang
  2. https://www.tensorflow.org/install/docker?hl=zh-cn , by tensorflow.
  3. https://huangwang.github.io/2018/10/09/Virtualenv简易教程/ , by jack huang
  4. https://laddiexu.github.io/tech/2017/11/04/TF-ODYourData.html , by 菁菁者莪
  5. https://blog.csdn.net/xunan003/article/details/78720189?utm_source=blogxgwz2, by xunan003