利用TensorFlow Object Detection API 训练自己的数据集

2023年10月20日 192次阅读 来源: shz0519

今天终于通过Tensorflow Object Detection API中的faster_rcnn_inception_resnet_v2来训练自己的数据了，参考：
数据准备
running pets
何之源的科普帖
简单记录如下：
这里，安装Tensorflow 和 Tensorflow Object Detection API就不说了.
为了保证训练过程中不修改原API源码目录结构，在其他地方新建了一个文件夹TF-OD-Test:如/home/shz/TF-OD-Test

1、训练数据集准备

在TF-OD-Test目录下新建文件夹data
由于数据集是PASCAL VOC2007格式的，并不符合API所需的TFRecord格式，需要转换，API已经为我们提供了代码. 但是这里我们需要做个简单的修改：
create_pascal_tf_record.py第160行 :

examples_path = os.path.join(data_dir, year, 'ImageSets', 'Main', 'aeroplane_' + FLAGS.set + '.txt')

改为：

examples_path = os.path.join(data_dir, year, 'ImageSets', 'Main',                           FLAGS.set + '.txt')

然后运行以下命令：

# From tensorflow/models/object_detection/
   python create_pascal_tf_record.py --data_dir=`自己的数据集路径` \
    --year=VOC2007 --set=train --output_path='TF-OD-Test/data/pascal_train.record'
   python create_pascal_tf_record.py --data_dir=/home/user/VOCdevkit \
    --year=VOC2007 --set=val --output_path='TF-OD-Test/data/pascal_val.record'

注意，文件data/pascal_label_map.pbtxt中的labels需要改为自己的。
执行后会在TF-OD-Test/data文件夹下生成pascal_train.record和pascal_val.record两个文件。
同时将修改后的pascal_label_map.pbtxt文件拷贝到TF-OD-Test/data目录下。

2、下载预训练模型

在TF-OD-Test目录下新建文件夹models

cd TF-OD-Test/models
wget http://storage.googleapis.com/download.tensorflow.org/models/object_detection/faster_rcnn_inception_resnet_v2_atrous_coco_11_06_2017.tar.gz
tar -xvf faster_rcnn_inception_resnet_v2_atrous_coco_11_06_2017.tar.gz
cp faster_rcnn_resnet101_coco_11_06_2017/model.ckpt.* TF-OD-Test/models

3、修改配置文件

将文件object_detection/samples/configs/faster_rcnn_inception_resnet_v2_atrous_pets.config复制到TF-OD-Test/models并打开做如下修改：
1）num_classes:修改为自己的classes num
2）将所有PATH_TO_BE_CONFIGURED的地方修改为自己之前设置的路径（共5处）

4、训练

终于可以训练了，进入object_detection目录，运行：

python train.py --train_dir='/home/shz/TF-OD-Test/train' --pipeline_config_path='/home/shz/TF-OD-Test/models/ssd_mobilenet/ssd_mobilenet_v1_pascal.config'

5、tensorboad：

输入命令：

tensorboard --logdir=/home/shz/TF-OD-Test/train

在浏览器中输入http://0.0.0.0:6006，就能看到训练曲线了

    原文作者：shz0519
    原文地址: https://zhuanlan.zhihu.com/p/27469690
    本文转自网络文章，转载此文章仅为分享知识，如有侵权，请联系博主进行删除。
点赞