TensorFlow Object Detection API 自行訓練模型教學,辨識特定物件

PASCAL VOC 2012 資料集

首先從 PASCAL VOC 2012 的網站上下載該資料集,整個資料集的大小大約是 2GB,下載時要稍微等一下:

cd models/research/

# 下載 PASCAL VOC 2012 資料集(2GB)
wget http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar

使用 create_pascal_tf_record.py 將 PASCAL VOC 2012 資料轉為 TFRecord 格式:

# 將 PASCAL VOC 2012 資料轉為 TFRecord 格式(training set)
python object_detection/dataset_tools/create_pascal_tf_record.py \
  --data_dir=VOCdevkit \
  --year=VOC2012 \
  --set=train \
  --label_map_path=object_detection/data/pascal_label_map.pbtxt \
  --output_path=object_detection/data/pascal_train.record

# 將 PASCAL VOC 2012 資料轉為 TFRecord 格式(validation set)
python object_detection/dataset_tools/create_pascal_tf_record.py \
  --data_dir=VOCdevkit \
  --year=VOC2012 \
  --set=val \
  --label_map_path=object_detection/data/pascal_label_map.pbtxt \
  --output_path=object_detection/data/pascal_val.record

這裡同樣以 Faster RCNN + Resnet101 + COCO 的模型作為預訓練模型,下載後放進 object_detection/data/ 目錄:

# Faster RCNN + Resnet101 以 COCO 資料集訓練好的模型
wget http://storage.googleapis.com/download.tensorflow.org/models/object_detection/faster_rcnn_resnet101_coco_11_06_2017.tar.gz

# 將模型參數放入 object_detection/data/
tar zxvf faster_rcnn_resnet101_coco_11_06_2017.tar.gz
cp faster_rcnn_resnet101_coco_11_06_2017/model.ckpt.* object_detection/data/

接下來要建立一個 .config 設定檔,這部分可參考 faster_rcnn_resnet101_voc07.config 這個內建的範例,將包含 PATH_TO_BE_CONFIGURED 的路徑修改一下,其餘的設定可以不需要更改,編輯完成後同樣儲存於 models/research/object_detection/data/ 之中,由於我們這裡使用的資料是 2012 年的,所以我把新的檔名變更為 faster_rcnn_resnet101_voc12.config

接著就可以開始訓練模型,訓練模型的方式跟之前相同:

# 設定檔路徑
PIPELINE_CONFIG="object_detection/data/faster_rcnn_resnet101_voc12.config"

# 訓練結果放置路徑
MY_MODEL_DIR="my_voc_model"

# 使用第一張 GPU 卡進行訓練
CUDA_VISIBLE_DEVICES=0 python object_detection/train.py \
  --logtostderr \
  --pipeline_config_path=${PIPELINE_CONFIG} \
  --train_dir=${MY_MODEL_DIR}/train

# 使用第二張 GPU 卡進行驗證
CUDA_VISIBLE_DEVICES=1 python object_detection/eval.py \
  --logtostderr \
  --pipeline_config_path=${PIPELINE_CONFIG} \
  --checkpoint_dir=${MY_MODEL_DIR}/train \
  --eval_dir=${MY_MODEL_DIR}/eval

# 啟動 TensorBoard
tensorboard --port=16006 --logdir=${MY_MODEL_DIR}

在訓練的過程中,同樣可以使用 TensorBoard 監看模型的狀況:

TensorBoard 畫面

驗證模型

模型訓練好之後,依照同樣的方式匯出模型:

# 設定檢查點檔案路徑
CHECKPOINT_NUMBER=800000
CKPT_PREFIX=${MY_MODEL_DIR}/train/model.ckpt-${CHECKPOINT_NUMBER}

# 將訓練好的模型匯出
python object_detection/export_inference_graph.py \
  --input_type image_tensor \
  --pipeline_config_path ${PIPELINE_CONFIG} \
  --trained_checkpoint_prefix ${CKPT_PREFIX} \
  --output_directory my_exported_graphs

並將新模型套用至前一篇教學的程式碼,跑出來的結果會像這樣:

測試結果

測試結果

參考資料:PythonProgramming.netLEARNING BY HACKINGPriya DwivediDat TranJiancheng Li

程式設計

5 留言

  1. 郭祐任

    請問我自己做的資料集辨識時會把同一種物件重覆框起是什麼問題

  2. guandeng

    请问训练好的模型只能识别狗和猫,那怎么和 COCO 資料集一起识别物体呢?

  3. Bruce

    請問為何你在 轉 TFRecord 檔後才產生2個檔案,而我卻產生了10幾個。
    情況和以下 github 一樣
    https://github.com/tensorflow/models/issues/4490
    而我也有改 face_only=False

    求您為我解惑~

  4. Henry

    您好,想請教一下在windows下訓練模型時,您提到的下列步驟:
    # 設定檔路徑
    PIPELINE_CONFIG=”object_detection/data/faster_rcnn_resnet101_pets.config”

    # 訓練結果放置路徑
    MY_MODEL_DIR=”my_model”
    是要如何執行呢?
    感謝QQ

  5. double

    請問假設訓練是辨識出狗和貓,那我拿大象的照片,可以得到無法辨識,這個答案嗎?

Leave a Reply