ディープラーニングモデルのオープンフォーマット「ONNX（Open Neural Network Exchange）」を使ってみる

ディープラーニング・アプリケーション開発キット”OpenVINO”を調べていた時、「ONNX」という単語に出会いました。
筆者の知らない単語だったので調べてみると、ONNX は「ディープラーニングモデルのオープンフォーマット」であることがわかりました。
今回はディープラーニングモデルのオープンフォーマット「ONNX」を紹介したいと思います。

ONNX って何？

ONNX（Open Neural Network Exchange）とは、異なるディープラーニングフレームワーク間でモデルの相互運用性を実現するプロジェクトのことです。
MicrosoftとFacebookが2017年に「ONNX」を発表しました。
現在では、多くの企業が参加し、AIの開発やデプロイの負担軽減に貢献するものとして期待されています。

ONNX登場以前はTensorflowやkeras, Pytorchなど、さまざまな形式のモデルをそれぞれのフレームワークで処理していましたが、ONNXモデルの登場で相互運用が可能になりました。
Tensorflowやkeras, Pytorchで作成したモデルをONNXモデルに変換したり、逆にONNXモデルを他のフレームワーク用のモデルに変換したりできます。

ONNX モデルを取得する

ONNX Model Zoo というサイトでONNXモデルが公開されているので、ここからONNXモデルをダウンロードしてみましょう。

今回は画像分類モデルの VGGモデルをダウンロードして試してみます。
モデルに VGG を選択したのは、シンプルなアーキテクチャと高い性能から現在も利用されるモデルであり、構造を理解するのに向いていると判断したからです。

VGG19-7

・
・その他多数の層
・

画像分類を試してみる

プログラムを jupyter lab で実行します。

ONNXランタイムがインストールされていない場合は、事前にランタイムをインストールしておきます。

pip install onnxruntime

GPU版の場合は

pip install onnxruntime-gpu

でインストールします。

カレントディレクトリに modelフォルダを作成し、その下に VGGモデル(vgg19-7.onnx) を配置します。

以下のコマンドを実行すると sess 変数に ONNXランタイムのセッションが作成されます。
以降はこのセッションを使って作業をします。

import onnxruntime
import numpy as np
onnx_model_path = "model/vgg19-7.onnx"
sess = onnxruntime.InferenceSession(onnx_model_path, providers=['CPUExecutionProvider'])

「onnxruntime.InferenceSession」メソッドの「providers」には「”CPUExecutionProvider”」を指定しました。
他に「”CUDAExecutionProvider”」を指定できますが、筆者の環境にはCUDA(Compute Unified Device Architecture)が無いので設定していません。

以下のコマンドを実行して、入力層を確認します。

input_name = sess.get_inputs()[0].name
print("Input name  :", input_name)
input_shape = sess.get_inputs()[0].shape
print("Input shape :", input_shape)

以下が出力されました。

Input name  : data
Input shape : [1, 3, 224, 224]

名前が「data」で、形式が「1, 3, 224, 224」であることがわかります。

以下のコマンドを実行して、出力層を確認します。

output_name = sess.get_outputs()[0].name
print("Output name  :", output_name)  
output_shape = sess.get_outputs()[0].shape
print("Output shape :", output_shape)

以下が出力されました。

Output name  : vgg0_dense2_fwd
Output shape : [1, 1000]

名前が「vgg0_dense2_fwd」で、形式が「1, 1000」であることがわかります。

以下のコマンドを実行して、分類する画像を読み込みます。

import cv2
import matplotlib.pyplot as plt

image = cv2.cvtColor(cv2.imread(filename="../data/image/cat3.jpg"), code=cv2.COLOR_BGR2RGB)
plt.imshow(image);

この画像は前回の mobilenet V3 モデルでうまく分類出来なかった画像です。
（我が家の愛猫「しおん」です。可愛くて仕方ありません）

以下のコマンドを実行して、画像をモデルに入力できる形に整形します。

resized_image = cv2.resize(src=image, dsize=(224, 224))
input_data = np.expand_dims(np.transpose(resized_image, (2, 0, 1)), 0).astype(np.float32)

以下のコマンド(run)を実行し、推論します。

out = sess.run([output_name], {input_name: input_data})

run コマンドの引数には出力層名称と、入力データを設定しています。
戻り地は出力層データです。

以下のコマンドを実行し、出力層のデータの最大値のインデックスを求めます。

result_index = np.argmax(out[0])

値を確認したところ「285」でした。

分類画像の名称インデックスのデータをここから取得します。
utilsフォルダの下に「synset.txt」として保存しました。

以下のコマンドを実行し、推論結果のインデックスに相当する文字列を出力します。

imagenet_classes = open("utils/synset.txt").read().splitlines()
imagenet_classes[result_index]

推論結果は、
n02124075 Egyptian cat
でした。（エジプト産の猫？）

前回の推論結果「ボーダー・コリー」(犬)よりは「猫」に寄ってきてくれてるので、良しとしましょう。

まとめ

学習済みONNXモデルを取得できれば、複数のフレームワークを使い分けることなくディープラーニングを用いた推論を行うことが出来ることがわかりました。
また機会があれば、Tensorflowや他のフレームワークのモデルをONNXモデルに変換して使ってみたいと思います。