本記事ではディープラーニングの領域で耳にすることの多いONNXについて、内容についてのまとめと検証を行なっていこうと思います。
本記事ではディープラーニングのモデルをONNX形式に変換した際の推論速度向上にフォーカスしており、Rasberry PiなどのエッジAIの領域で推論速度の要件を満たすための選択肢の1つとして、読者の方の参考になればと思います。
ONNXについて
ONNXはOpen Neural Network eXchangeの略で、DeepLearningモデルを様々なフレームワーク間で交換するためのフォーマットです。読み方はオニキスです。
Tensorflowを用いて学習させたモデルをpytorchで推論させる際にモデルをそのまま用いることはできませんが、モデルをONNX形式に変換すれば可能となります。
今回の記事では、このフォーマットの互換性ではなく、ONNX形式でモデルを推論させた際の推論速度の向上と、副次作用としての精度の劣化について主にまとめます。
ONNXで速度向上が起きる理由と有用場面
ディープラーニングモデルの速度向上というと、枝刈り、量子化、蒸留などを耳にすることがありますが、ONNX化がもたらす推論速度向上はそれらとは異なります。
ONNXの最適化は計算グラフの中で推論時には冗長になる部分をよりシンプルな形に置き換えるといった最適化処理を行うもので、例えばBatchNormalizationをConvolutionに還元したり、一度だけの計算で済む処理を別のグラフに切り出して計算回数を削減したりといったことを行います。
推論をGPU等で行う場合や処理件数が少ない場合など、処理時間がシステムにとってネックになりづらい際は速度向上のメリットよりもONNX化に必要なワークロードや精度が低下する可能性があるというデメリットの方が目立ってしまいますが、GPUのないPCやRasberry Piといった機器で推論処理を行いたい場合、推論にかかる時間が無視できない場合があります。そういった場合に推論速度を上昇できるのはシステムの実用性に与えるメリットも大きくなります。
ONNX化に必要なワークロードや精度が低下する可能性についても、以降で示す手順や検証結果からそれほど大きなデメリットではないことがお伝えできればと思います。
環境
- MacBook Air
- プロセッサ数：1
- コア数：２
- メモリ：8GB
- Python3.7.6
- Pytorch 1.8.1
- onnxruntime 1.7.0
速度・精度検証
今回はpytorchを用いて学習させた画像認識モデル（ResNet50）に対して、ONNX形式に変換する前と後で推論速度と精度の比較を行います。
フレームワークは学習・推論共に Pytorch を用いています。
学習条件としてepoch数を20と50にした2種類のモデルで検証しました。
モデルの学習と推論として鳥の種類の分類器を用いました。
学習・テスト画像の収集として、pythonのicrawlerを利用してブラウザから自動で画像収集を行いましたが、クイックに画像を集めてデータセットを作成する際に結構便利と感じました。
たとえば犬の画像を100枚収集したい場合は以下の様に実行することで、birdディレクトリの中に「犬」というキーワードの検索結果画像を保存していくことができます。
from icrawler.builtin import BingImageCrawler
crawler = BingImageCrawler(storage={"root_dir": "bird"})
crawler.crawl(keyword="犬", max_num=100)
今回の検証では鳥の種類以下9種類を上記手順で収集して画像認識を行いました。
［ハッカチョウ、ヨシキリ、タイヨウチョウ、ササゴイ、カラス、アカゲラ、バン、ウミワシ、カモメ］
9種で合計190枚の画像を収集し、145枚の学習画像と45枚のテスト画像に分けました。
ResNet50モデルででそれぞれpthファイルのモデルを作成し、ONNXファイルへの変換を行います。
pthファイルからonnxファイルへの変換は以下のコードで行います。Pytorch自体にモデルをONNXに変換してエキスポートする機能が備わっています。
import torch
import torch.nn as nn
from torchvision import models
# Load class names
with open('image_classes.txt') as f:
class_names = [line.strip() for line in f.readlines()]
# Define model
model = models.resnet50(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, len(class_names))
# Load model
model.load_state_dict(torch.load(PATH,map_location=torch.device('cpu')))
# Convert pth model to onnx
batch_dummy = torch.zeros(1,3,224,224)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()
torch.onnx.export(model, batch_dummy, outonnxpath, opset_version=11)
モデルの変換の際には学習の際に用いた推論クラスのリストを与えてあげる必要があるので、今回は image_classes.txt を用意して読み込んでいます。
このリストは推論の際に用います。
onnxモデルによる画像推論は以下のコードで行います。
onnxruntimeモジュールがPython3.7移行でしかダウンロードできない都合上python3.6以下の環境では下記コードは実行できません。
import onnxruntime
import torch
from PIL import Image
from torchvision import transforms
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
data_transforms = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# Load class names
with open('image_classes.txt') as f:
class_names = [line.strip() for line in f.readlines()]
# Prepare test image
img = Image.open(testimage)
img_t = data_transforms(img)
batch_t = torch.unsqueeze(img_t, 0)
out = batch_t.to(device).detach().numpy().copy()
# Inference with ONNX model
sess = onnxruntime.InferenceSession(MODELPATH)
input_name=sess.get_inputs()[0].name
pred_onnx=sess.run([], {input_name: out})
out_onnx = torch.from_numpy(pred_onnx[0][0]).clone()
_, index = torch.max(out_onnx, 0)
predclass = class_names[index]
MODELPATH は別途定義して与える必要があります。
検証結果
推論速度は推論した枚数（45枚）を推論時間で除したあ値で、3回テストを行った平均としています。
- epoch20で学習したモデル
| Test Accuracy | Inference Speed (枚 / sec) | |
|---|---|---|
| pthモデル | 0.978 | 2.66 |
| onnxモデル | 0.978 | 9.00 |
※どちらのモデルもテスト画像45枚のうち1枚間違えましたが、同じ画像で間違えていました。
- epoch50で学習したモデル
| Test Accuracy | Inference Speed (枚 / sec) | |
|---|---|---|
| pthモデル | 1.0 | 2.45 |
| onnxモデル | 1.0 | 9.00 |
どちらの学習条件でも、ONNX変換したことによる推論精度の低下は見られず、推論速度は3倍以上向上していました。
1秒で5枚の画像を推論処理したいという要件がAIに求められる場合、そのままでは実現できませんが、ONNX化することで容易に実現できるようになります。
ただし、今回の検証では推論精度の低下はありませんが、必ずそのようになるわけではなく、モデルのタスクが画像分類から物体検出やセグメンテーションとなったり処理したい画像の複雑さによっては精度の低下が見られる可能性もあります。今回の検証では再現できませんが、実際に私がONNXを使用していてそのような場面に遭遇したことも少ないですがありました。精度の低下が見られるかどうか、見られる場合にどの程度低下するかは今回のようにテスト画像での検証を行なって確認することを推奨します。その結果精度低下が起きない、もしくは起きたとしてもそれがシステム上要件の範囲に収まると考えられるのであれば推論速度向上とのバランスを鑑みて採用するとよいと考えられます。
今回の検証はMacBookのCPUのみで、モデルも画像認識モデルのみだったので、今後Rasberry Piでの検証や物体検出といったタスクでも引き続き検証して追記していきたいと思います。