Tensorflow Object Detection 训练自己数据

本文提供给大家参考，若有不对的地方还请各位指教。

ok，学习Tensorflow物体检测之后，我们发现初始训练数据并不是我们自己提供的，那怎么才能换成我们自己的训练集呢？

下面是直接采用object-detection model 测试的结果，显示出了鼠标和杯子。下一步我们想加一些自己的数据给参与训练。

测试图片.png

准备数据集：

我们可以自己生成TFRecord数据文件，推荐工具LabelImg。现在我们使用PASCAL VOC 2012数据集为例。下载地址

下载之后我们通过object-detection model 提供的create_pascal_tf_record.py文件将VOC2012数据转化为TFRecord文件。
在create_pascal_tf_record.py文件中修改以下地方：

python create_pascal_tf_record.py \
    --label_map_path=data/pascal_label_map.pbtxt \
    --data_dir=data/VOCdevkit \ 
    --year=VOC2012 \ 
    --set=train \
    --output_path=data/pascal_train.record \ 
python create_pascal_tf_record.py \
    --label_map_path=data/pascal_label_map.pbtxt \
    --data_dir=data/VOCdevkit \
    --year=VOC2012 
    --set=val \
    --output_path=data/pascal_val.record

train 和 val 分别修改生成 pascal_train.record以及pascal_val.record文件。

数据.png

下载模型

官方提供了不少预训练模型。我是地址
我推荐下载ssd_mobilenet_v1_coco模型，比较轻量级，之后测试训练速度相比其他比较快。

因为下载的是ssd_mobilenet_v1_coco模型，所以把object_detection/samples/configs/ssd_mobilenet_v1_pets.config文件复制到一个新建文件夹中，目录最好是在object-detection文件下。把之前下载的模型也复制到这个文件夹中，对
ssd_mobilenet_v1_pets.config进行修改：

num_classes：20
fine_tune_checkpoint：模型model.ckpt路径
train_input_reader：{
   input_path：pascal_train.record路径
   label_map_path：../data/pascal_label_map.pbtxt
}
eval_input_reader：{
   input_path：pascal_val.record路径
   label_map_path：../data/pascal_label_map.pbtxt
}

开始训练

配置object_detection/train.py文件（可以复制新建一份做出修改）：

python ../../train.py \
    --logtostderr \
    --pipeline_config_path=ssd_mobilenet_v1.config 的路径
    --train_dir=train_logs（随意一个路径，保存模型使用）

在object_detection train.py文件中，我们还需要把下载的model文件，内容可能包括checkpoint、frozen_inference_graph.pb等文件一共放置于pipeline_config_path同目录下。

运行train.py文件，ok已经开始打印日志文件了：

训练.png

训练完成之后调用export_inference_graph.py把训练好的数据生成model就可以进行检测。

python export_inference_graph.py 
--input_type image_tensor
--pipeline_config_path : 最开始复制的config文件，本文中举例的是ssd_mobilenet_v1_pets.config
--trained_checkpoint_prefix ：训练好的mode.ckpt文件 ，例如 model.ckpt-n 
--output_directory : 自己想保存的文件路径

执行python train.py报错：trainer.py“ModuleNotFoundError: No module named 'deployment'”，或者报错No module named 'nets'
错误原因，这两个model都是在TensorFlow/models/slim中，需要把models/slim加入到PYTHONPATH环境变量中。