搭建 MobileNet-SSD 开发环境并使用 VOC 数据集训练 TensorFlow 模型

原文地址:搭建 MobileNet-SSD 开发环境并使用 VOC 数据集训练 TensorFlow 模型

0x00 环境

OS: Ubuntu 1810 x64
Anaconda: 4.6.12
Python: 3.6.8
TensorFlow: 1.13.1
OpenCV: 3.4.1

0x01 基础环境配置

Anaconda 下载地址: Anaconda-4.6.12-Linux

本文中安装位置为 /usr/local/anaconda3

修改默认的 python 版本为 3.6

conda install python=3.6

安装 OpenCV 3.4.1

conda install opencv=3.4.1

安装 TensorFlow 1.13.1

conda install tensorflow=1.13.1

0x02 TensorFlow Models

下载地址: Github - TensorFlow Models

下载后得到一个 models-master.zip 文件,解压后移动到 /usr/local/anaconda3/lib/python3.6/site-packages/tensorflow文件夹下,并重命名为 models

unzip models-master.zip
mv models /usr/local/anaconda3/lib/python3.6/site-packages/tensorflow

进入 models/research/目录,并编译 protobuf

cd /usr/local/anaconda3/lib/python3.6/site-packages/tensorflow/models/research
protoc object_detection/protos/*.proto --python_out=.

安装 object_detection 库

python setup.py build
python setup.py install

设置 PYTHONPATH

export PYTHONPATH=$PYTHONPATH:/usr/local/anaconda3/lib/python3.6/site-packages/tensorflow/models/research
export PYTHONPATH=$PYTHONPATH:/usr/local/anaconda3/lib/python3.6/site-packages/tensorflow/models/research/slim

直接执行以上命令只会在当前终端生效,将以上命令写入 ~/.bashrc并执行如下命令可以永久保存

source ~/.bashrc

测试 object_detection 库是否安装成功

python object_detection/builders/model_builder_test.py

进入 object_detection/目录并启动 jupyter-notebook,测试目标检测

cd object_detection/
jupyter-notebook

在浏览器中打开 http://localhost:8888/,进入 jupyter-notebook 控制台,打开 object_detection_tutorial.ipynb 文件并执行,待模型下载完成并检测完成后会在页面底部出现两张标注后的图片

0x03 训练

下载 VOC 2012 数据集: VOCtrainval_11-May-2012.tar

object_detection/目录下创建目录 ssd_model,并解压数据集至 object_detection/ssd_model

mkdir ssd_model/
cd ssd_model
tar xvf VOCtrainval_11-May-2012.tar

返回 research/目录,执行 train 和 val 脚本

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

这两个脚本会在 ssd_model/目录下生成 pascal_train.record 和 pascal_val.record 两个文件,各 600M 左右

复制配置文件,在此基础上修改,并训练数据

cp object_detection/data/pascal_label_map.pbtxt object_detection/ssd_model/
cp object_detection/samples/configs/ssd_mobilenet_v1_pets.config object_detection/ssd_model/

pascal_label_map.pbtxt 文件中保存了数据集中有哪些 label

将 ssd_mobilenet_v1_pets.config 中的 num_classes 改为 pascal_label_map.pbtxt 中列出的文件数量,这里是 20,并修改迭代次数 num_steps,并将配置文件末尾的路径按照如下格式修改

train_input_reader: {
  tf_record_input_reader {
    input_path: "/usr/local/anaconda3/lib/python3.6/site-packages/tensorflow/models/research/object_detection/ssd_model/pascal_train.record"
  }
  label_map_path: "/usr/local/anaconda3/lib/python3.6/site-packages/tensorflow/models/research/object_detection/ssd_model/pascal_label_map.pbtxt"
}

eval_input_reader: {
  tf_record_input_reader {
    input_path: "/usr/local/anaconda3/lib/python3.6/site-packages/tensorflow/models/research/object_detection/ssd_model/pascal_val.record"
  }
  label_map_path: "/usr/local/anaconda3/lib/python3.6/site-packages/tensorflow/models/research/object_detection/ssd_model/pascal_label_map.pbtxt"
  shuffle: false
  num_readers: 1
}

下载 ssd_mobilenet 至 ssd_model/目录下,解压并重命名为 ssd_mobilenet

ssd_mobilenet: ssd_mobilenet_v1_coco_11_06_2017.tar.gz

tar zxvf ssd_mobilenet_v1_coco_11_06_2017.tar.gz
mv ssd_mobilenet_v1_coco_11_06_2017 ssd_mobilenet

将 ssd_mobilenet_v1_pets.config 中 fine_tune_checkpoint 修改为如下格式的路径

fine_tune_checkpoint: "/usr/local/anaconda3/lib/python3.6/site-packages/tensorflow/models/research/object_detection/ssd_model/ssd_mobilenet/model.ckpt"

使用 train.py 脚本训练模型

注意:脚本可能位于 object_detection/object_detection/legacy/目录下

这里位于 object_detection/legacy/目录

python ./object_detection/legacy/train.py --train_dir ./object_detection/legacy/train/ --pipeline_config_path ./object_detection/ssd_model/ssd_mobilenet_v1_pets.config

运行 export_inference_graph.py 脚本将训练出的模型固化成 TensorFlow 的 .pb 模型,其中 trained_checkpoint_prefix 要设置成 model.ckpt-[step],其中 step 要与训练迭代次数相同

python ./object_detection/export_inference_graph.py --input_type image_tensor --pipeline_config_path ./object_detection/ssd_model/ssd_mobilenet_v1_pets.config --trained_checkpoint_prefix ./object_detection/legacy/train/model.ckpt-9000 --output_directory ./object_detection/ssd_model/model/

转换后生成的 .pb 模型位于 object_detection/ssd_model/model/目录下

将 pascal_label_map.pbtxt 作为 label 文件

展开阅读全文

没有更多推荐了,返回首页