ONNXでモデルを表現しよう

ONNX（Open Neural Network Exchange） は機械学習やディープラーニングなどのAIモデルを表現するためのオープンソースのフォーマットです。プラットフォームに依存しない、という特徴があり、ONNXモデルを経由することによりPyTorchで学習したモデルをTensorflowで推論することや、ワークステーション上のPythonスクリプトで学習したモデルをエッジデバイスのC++アプリで推論することなどが可能になります。

Raspberry Pi 5のAI KitであるHailo-8がONNXに対応しているとのことなので、ゴールはUbuntu24.04 + RTX3060Tiでモデルを学習させて、Raspberry Pi 5によりエッジで推論する構成を目指しているのですが、本記事では、まずMacBook Pro上でONNXを味見してみようということを目的に、遊んでみました。

関連記事

Netron

なお、本記事ではネットワークが変換できているかを確認するために、Netronを利用します。Netronはモデルを可視化できるツールで、インストールすることなくWebから利用可能です。Netronを使うことにより深層学習の各層の結合を確認できる他、各層の入力テンソルも確認することができます。下記の例は入力テンソルが (1, 3, 224, 224) となっているResNetの例です。

学習済みのResNetをONNXモデルとして扱う

PyTorchのモデルをONNXモデルへ書き出す

まずPyTorchのtorchvisionにバンドルされている学習済みのresnet101をONNXモデルに変換してみましょう。resnet101は入力に1枚の画像、RGBチャネル、224x224ピクセルの画像をとり、出力に1000種類のクラスの確率をとるモデルです。

resnet101をONNX形式へ変換するmodel_converter.pyの内容は以下となります。モデルをtorchvision.modelsから読み込みresnet変数に格納した後、推論モードへ切り替え、ダミーの入力データ(1, 3, 224, 224)を準備し、onnx.exportを呼び出すことでONNXモデルのファイルを生成できます。

# 実行のベースとなるtorchをインポート
import torch
# torchvisionからモデルと画像変換機能をインポート
from torchvision import models
# PyTorchのONNX拡張
import torch.onnx as onnx

##############
# resnetの準備
##############
print(dir(models))
# モデルに含まれているresnetを、学習済みの状態で取得
resnet = models.resnet101(pretrained=True)
# resnetのネットワーク構造を表示
print(resnet)
# resnetを推論モードに切り替え
resnet.eval()

# ダミーデータを作成する
input_image = torch.zeros((1, 3, 224, 224))
# ONNX形式で書き出す
onnx.export(resnet, input_image, "model.onnx")
# 終了する
quit()

PyTorchのモデルからONNXモデルへの変換手順は以下の通りです。変換にはonnxとonnxruntimeをインストールする必要があるため、実行前に仮想環境にpipコマンドでインストールします。

# 作業ディレクトリを作成する
$ mkdir pytorchbox
$ cd pytorchbox 

# 仮想環境を作成する
$ conda create --prefix ./env python=3.8
# 仮想環境を有効化する
$ conda activate ./env
# PyTorchをインストールする
(env) $ conda install pytorch::pytorch torchvision torchaudio -c pytorch
# onnxとonnxランタイムをインストールする
(env) $ pip install onnx
(env) $ pip install onnxruntime

# PyTorchモデルをONNXへ変換するプログラム
(env) $ ls *.py
model_converter.py

# 変換を実行する
(env) $ python model_converter.py
...

# ONNXモデルができていることを確認する
(env) $ ls *.onnx
model.onnx

生成した model.onnx をNetronから開くことで、下記のようにネットワークを可視化可能です。入力テンソルが(1, 3, 224, 224)で、出力テンソルが(1, 1000)となっていることを確認することができますね。

ONNXモデルをPythonで読み込んで推論する

続いて、ONNXモデルをPythonで読み込んで推論をしてみましょう。基本的な構造は下記記事と同等で、変更されている箇所は、推論の呼び出し方や後処理周りです。

利用したパッケージ

ONNXモデルの実行にPyTorchは必要ありませんが、今回はカメラからの入力画像やテンソルを加工しやすいことからtorch、torchvision、PILをインポートしています。また、ONNXモデルを扱うためonnxとonnxruntimeをインポートしています。

全体の流れ

ONNXモデルはort.InferenceSessionにより読み込みます。モデルの入力層と出力層の名前を取得して、推論を実行する際(ort_session.run)に指定します。推論の入力となるWEBカメラで撮影した画像をPyTorchにより(1, 3, 224, 224)のテンソルに変換して.numpy().copy()でnumpy配列とします。これをort_session.runの引数に与えて推論を実行します。推論結果は1000クラスそれぞれの確率となりますので、確率を昇順に並べ、カメラに何が映ったかを判別します。今回は、出力層をNetronを見ると最後のGemm層で(x2048)していましたので、2048で割ってパーセンテージに変換し、50%以上であれば出力するようにしました。

# 実行のベースとなるtorchをインポート
import torch
# torchvisionからモデルと画像変換機能をインポート
from torchvision import transforms
# 画像を扱うライブラリPILをインポート
from PIL import Image

# ONNXを制御するパッケージをインポート
import onnx
import onnxruntime as ort
from onnxruntime import GraphOptimizationLevel, InferenceSession, SessionOptions
# 数値処理用のパッケージをインポート
import numpy as np
# Webカメラを制御するOpenCVをインポート
import cv2

# /dev/video0を指定
DEV_ID = 0
# 撮影するサイズ
WIDTH = 640
HEIGHT = 480
# 認識の閾値
THRESHOLD = 50

##############
# resnetの準備
##############
# ONNX実行時のオプションを設定する
options = SessionOptions()
options.intra_op_num_threads = 1
options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL

# ONNXモデルを読み込む
ort_session = ort.InferenceSession('model.onnx')
ort_session.disable_fallback()
input_name = ort_session.get_inputs()[0].name
output_name = ort_session.get_outputs()[0].name

# 入力層の名前を取得
print("input:")
for session_input in ort_session.get_inputs():
    print(session_input.name, session_input.shape)
# 出力層の名前を取得
print("output:")
for session_output in ort_session.get_outputs():
    print(session_output.name, session_output.shape)

# resnetの分類(クラス)をlabelsに読み込む
with open('./imagenet_classes.txt') as f:
    labels = [line.strip() for line in f.readlines()] 

# OpenCVで撮影した画像をresnetの入力テンソルに変換する
preprocess = transforms.Compose([
    transforms.ToPILImage(), # numpy.ndarray -> pillow
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

##############
# Webカメラに接続する
##############
cap = cv2.VideoCapture(DEV_ID)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, WIDTH)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, HEIGHT)
cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)

##############
# OpenCVで撮影し、resnetテンソルに変換
##############
while(True):
    try:
        # 1フレーム撮影する
        ret, frame = cap.read()
        if ret:
            # 撮影したフレームを表示
            cv2.imshow('frame', frame)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
            # 画像をBGRからRGBの並びに変換
            input_image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            # 画像をresnetの入力テンソルに変換する
            input_tensor = preprocess(input_image)
            # 画像をバッチに内包する(1次元追加)
            batch_t = torch.unsqueeze(input_tensor, dim=0)
            # ONNXに渡すデータへと変換する
            input_onnx = batch_t.numpy().copy()

            ##############
            # 入力テンソルをresnetで推論、結果を表示する
            ##############
            # カメラの画像をresnetで推論
            # runの第一引数：出力層の名前
            # runの第二引数：入力narray
            out_onnx = ort_session.run([output_name], {input_name: input_onnx})

            #out = torch.from_numpy(out_onnx[0][0]).clone()
            #_, index = torch.max(out, 0)
            out = out_onnx[0][0]
            # 確率を昇順にソートする
            sorted = np.sort(out)[::-1]
            sorted_idx = np.argsort(out)[::-1]
            # パーセンテージに変換
            parcentage = (sorted * 100 / 2048) * 100 
            # 先頭のアイテムを取り出す
            main_percentage = parcentage[0]
            main_label = labels[sorted_idx[0]]
            #print('%s(%s)' % (main_label, "{:.1f}".format(main_percentage, 1)))
            
            # 確率一定以上であれば結果を出力
            if main_percentage > THRESHOLD:
                print('%s(%s)' % (main_label, "{:.1f}".format(main_percentage, 1)))

    except KeyboardInterrupt:
        print('abort')

スクリプトの実行手順は以下の通りです。

# 作業ディレクトリへ移動
(base) $ cd pytorchbox 
(base) $ pwd
/Users/shino/pytorchbox

# 実行にはスクリプトの他に
# クラス名のリスト(imagenet_classes.txt)と
# ONNXモデルファイル(model.onnx:前手順でエクスポートしたもの)が必要となる
(base) $ ls
fetch_model_and_detect.py
model.onnx
env
imagenet_classes.txt
model_converter.py

# 仮想環境を有効にする
(base) $ conda activate ./env 
(env) $ 

# 実行
(env) $ python fetch_model_and_detect.py 
# book jacket(58.6)
# book jacket(56.5)
# book jacket(60.6)
# ... ターミナルを閉じて終了

認識できましたね！これでまた一歩前進です！

参考させていただいた記事

今回も多数の記事を参考にさせていただきました。ありがとうございます。