tensorflow-onnx｜GitHubを参考にColaboratoryを使って、 TensorFlowの学習済みモデルをONNXモデルに変換します
以下のコマンドでインストールします
!pip install --user -U tf2onnx
前回、ONNX RuntimeとYoloV3でリアルタイム物体検出｜はやぶさの技術ノートについて書きました
今回は『SSDでリアルタイム物体検出』を実践します
ONNXモデルをエクスポートできる深層学習フレームワークは複数ありますが、
SSD系の学習済みモデルについては、 Tensorflow detection model zooが非常に充実してます。
なので、TensorFlowモデルをONNXモデルに変換して、最終的にONNXRuntimeで物体検出（推論）を行います
最初に、パスやファイル名・フォルダ名などを定義します
import os
import sys
ROOT = os.getcwd()
WORK = os.path.join(ROOT, "work")
MODEL = "ssdlite_mobilenet_v2_coco_2018_05_09"
os.makedirs(WORK, exist_ok=True)
# force tf2onnx to cpu
os.environ['CUDA_VISIBLE_DEVICES'] = "0"
os.environ['MODEL'] = MODEL
os.environ['WORK'] = WORK
Tensorflow detection model zoo の中から好きなモデルを選びます
個人的に、リアルタイム物体検出が好きなので、”軽快に動作する”ssdlite_mobilenet_v2_cocoを採用します
以下のコマンドでダウンロードします
!cd $WORK; wget -q http://download.tensorflow.org/models/object_detection/$MODEL.tar.gz
!cd $WORK; tar zxvf $MODEL.tar.gz
ダウンロードしたファイルを確認
!ls $WORK/$MODEL
以下のコマンドで【input】と【output name】の定義を確認できます
!saved_model_cli show --dir $WORK/$MODEL/saved_model/ --tag_set serve --signature_def serving_default
inputについては、以下の通りです
The input is a batch of images in NHWC format and the used data type it uint8.
学習済みモデルが”frozen_inference_graph.pb”と”saved_model”で保存してあるので、それぞれONNXモデルに変換する方法を紹介します
tf2onnxを使って、TensorFlowモデル”saved_model”をONNXモデルに変換
!python -m tf2onnx.convert --opset 10 --fold_const --saved-model $WORK/$MODEL/saved_model --output $WORK/$MODEL.onnx
同様に、tf2onnxを使って、TensorFlowモデル”frozen_inference_graph.pb”をONNXモデルに変換
!python -m tf2onnx.convert --graphdef $WORK/$MODEL/frozen_inference_graph.pb --output $WORK/$MODEL.frozen.onnx \
--fold_const --opset 10 \
--inputs image_tensor:0 \
--outputs num_detections:0,detection_boxes:0,detection_scores:0,detection_classes:0
変換したONNXモデルを推論用ディバイスで使用するため、「Googleドライブに保存」⇒「ダウンロード」を行います
まずは、【秒速で無料GPUを使う】深層学習実践Tips on Colaboratory を参考に、以下の手順でONNXモデルをGoogleドライブに保存します
以下のコードでGoogleドライブをマウントします
from google.colab import drive
drive.mount('/content/drive')
正常にマウントできれば、以下のコマンドで対象Google Driveの内容が表示されます
!ls /content/drive/'My Drive'
以下のコマンドでONNXモデルをGooglドライブにコピーします
!cp $WORK/$MODEL.onnx /content/drive/'My Drive'
改めて、対象Google Driveの内容を表示させるとssdlite_mobilenet_v2_coco_2018_05_09.onnxが保存を確認できます
!ls /content/drive/'My Drive'
今回は、”ssdlite_mobilenet_v2_coco_2018_05_09.onnx”を取得しましたが、
同様の手順で他のONNXモデルも取得できます