搭建 MobileNet-SSD 開發環境並使用 VOC 資料集訓練 TensorFlow 模型

0x00 環境

OS: Ubuntu 18.10 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 檔案