linux自定义model,使用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驱动即可。

安装docker

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

安装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

下载最新支持GPU和Python3的tensorflow映像[2]

1docker pull tensorflow/tensorflow:latest-gpu-py3

测试tensorflow映像

1

2docker 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

9git 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

12item {

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

9workon 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

12workon 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 API1

2

3

4git clone https://github.com/cocodataset/cocoapi.git

cd cocoapi/PythonAPI

make

cp -r pycocotools /models/research/

编译Protobuf1

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=.

添加库到PYTHONPATH1

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

模型训练

数据准备将label_map文件、训练用TFRecord文件、测试用TFRecord文件复制到tensorflow/models/research/object_detection/data文件夹下。1

2

3

4

5# 推荐的数据目录结构

+ data

- geometry.pbtxt

- geometry_test.record

- geometry_train.record

下载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

修改解压后模型文件夹中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.pbtxt1

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。

启动docker1

2# 启动docker，并使用-v参数将目录/home/$USER挂载到容器上

docker run --runtime=nvidia -v /home/$USER:/home/$USER -it tensorflow/tensorflow:latest-gpu-py3 bash

执行训练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

使用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 步后自动保存模型。

参考文献