株式会社神戸デジタル・ラボ DataIntelligenceチームの原口です。
今回はAIモデルをCPUで高速に実行できると噂のONNX・OpenVINOを用いてYOLOv5の高速化に取り組みます！
AIモデルをスマートフォンやIoTデバイスで動作させたい全国のエンジニア必見です！
はじめに
皆さんはAIを実行している際
「このモデルすごいけど、めっちゃ遅くない？」
「GPUがないマシンだとなんもできんから実質案件で使えんやん・・・」
などと考えたことはありませんか？私は日々精度と速度のトレードオフに悩まされています。
「昔のモデルもCPU推論遅かったっけ？」と考えたそこのあなた！
今のモデルと昔のモデルで計算量がどれくらい増えたかご存知ですか？少し比較してみましょう。
| モデル名 | 発表年度 | 計算量 |
|---|---|---|
| VGG16 | 2014 | 15.4GFLOPS |
| Swin-L | 2021 | 103.9GFLOPS |
VGG16は畳み込みニューラルネットワーク（CNN）の名を世に知らしめた有名なモデルですね。このモデルが発表されたのが2014年ですね。 もう一方のSwin-Lはコンピュータビジョン界隈に激震を走らせたViTの後継モデルです。このモデルが発表されたのが2021年です。
計算量（FLOPS）は推論結果を出すまでの浮動小数点演算の計算回数を表しています。1FLOPSは1回の計算です。 VGG16が約150億回の計算をして推論結果を出力していたのに対し、Swin-Lでは約1兆回の計算をしないと推論結果が出力されません。約6倍の計算回数が必要となります。それだけ差があれば遅く感じるのも当然ですね。
推論速度は昔のモデル程度で、精度は今のモデルというAIが欲しいですがそうは問屋が卸しません。
そうなってくるとすることは決まってきます。そう、モデルの高速化です。
AIモデルの高速化には様々な手法がありますが、今回の実験ではYOLOv5モデルをONNX・OpenVINOに変換することで高速化を実現します！
ONNXとは
ONNX（Open Neural Network Exchange）とは様々なAIモデルを様々なAIライブラリで利用できるように変換できるライブラリで、言語界の英語のような存在です。
一般的に機械学習にはTensorFlowやPyTorchなどのAIライブラリを利用します。学習によって作成されたAIモデルは各AIライブラリに依存した形式で保存されます。そのためTensorFlowで作成されたAIモデルはPyTorchで利用できません。その逆も然りです。
そのため「お、このモデルちょっと良さそうやん。使って・・・、うわ使ってる環境と違うやん。しかも依存関係の指定めちゃくちゃあるやん。あー・・・。」のようなことが頻発していました。
このような場合にONNXは非常に有用です。上で説明したように、ONNXは様々な機械学習ライブラリの学習済みモデルをONNX形式へ変換することができます。つまりONNXが実行できる環境が一度準備できれば、以降AIモデルの実行のために複雑な環境構築をする必要がなく、様々なモデルをONNXへ変換すれば実行できるようになります。
またONNXへ変換することでCPU上での高速推論も可能になります。TensorFlowやPyTorchなどの機械学習ライブラリはモデルの学習・推論などあらゆる処理をカバーしています。そのためCPUでの推論に特化しておらず、推論時に必要のない処理が実行されていることがあります。
ONNXは推論のみを行うライブラリであり、不要な処理が取り除かれているためCPUでの高速推論が実現できるというわけです。
実験環境
実験に移る前に今回の実験環境について簡単にご紹介します。
- OS：Windows10
- CPU：Intel Core i7-10510U@1.80GHz
CPU推論はCPUの性能によって速度が変化します。今回の実験は上記環境で行っていることをご確認ください。
PyTorch vs ONNX
ではYOLOv5を用いてPyTorchとONNXの実行速度の比較を行ってみましょう。
まずはPyTorchの速度検証から始めます。
PyTorchの速度検証
実験に移る前に物体検出モデルの1種である「YOLOv5」の準備をしましょう。
YOLOv5のインストールはたった一行で済みます。本当にすごいですよね・・・。
pip install -qr https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt
続いて速度検証を行います。
PyTorchでの推論は次のコードで実行しました。
from time import time import cv2 import numpy as np import torch # バウンディングボックスの描画関数 def plot_bbox(frame, bboxes, time): for bbox in bboxes: bbox = bbox.numpy().astype(int) frame = cv2.rectangle(frame, bbox[0:2], bbox[2:4], (0, 255, 0), 3) cv2.putText(frame, text=f"time {time}", org=(0, 40), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1.0, color=(0, 255, 0)) return frame def main(): # カメラ接続 cap = cv2.VideoCapture(0) print("cam ok") # モデル読み込み。今回は最も軽いyolov5n6を利用 model = torch.hub.load('ultralytics/yolov5', 'yolov5n6', pretrained=True) print("model ok") # time list 30 frameの平均値を測る為 time_list = [0 for _ in range(30)] pointer = 0 while True: # 動画の1フレームを取得 ret, frame = cap.read() # 計測開始 start = time() # YOLOv5による推論 output = model(frame[:, :, ::-1]) # 30 fraemの平均を出すために保存 time_list[pointer % 30] = time() - start pointer += 1 # 描画関数 frame = plot_bbox(frame, output.xyxy[0], np.mean(time_list)) cv2.imshow("test", frame) cv2.waitKey(1) if __name__ == "__main__": main()
上記コードを実行すると新しい画面が表示され、そこに検出結果が表示されます。検出結果の左上に掛かった時間が表示されています。
・・・150ms程度ですね。テレビの描画速度が33msなので5倍の時間がかかっていますね。
ではONNXを用いて高速化に取り組みましょう！
ONNXの速度検証
続いてONNXでの推論を実行しましょう。ONNXで実行するためにはPyTorchモデルをONNXモデルへ変更する必要があります。
まずはONNXへの変換に取り組みます。
ONNXへの変換は次のコードで実行可能です。
import torch import torch.onnx as torch_onnx def main(): # PyTorchモデルの呼び出し model = torch.hub.load('ultralytics/yolov5', 'yolov5n6', pretrained=True) # ダミーインプットの作成 dummy_input = torch.randn(1, 3, 512, 640) # ONNXへの変換 output = torch_onnx.export(model, dummy_input, "./yolov5n6_onnx.onnx", input_names=["input"], output_names=["output"]) if __name__ == "__main__": main()
ONNXへの変換にはダミーインプット（適当なデータ。なんでも良い。）を要します。理由はPyTorchのモデル設計がDefine by Runであるからです。Define by Runタイプでは、モデルの構造をデータを流しながら確定させていきます。
つまり、入力データによって中間のデータの持ち方が変化するわけです。ONNXへ変換するにはモデル構造を確定させる必要があります。そこでダミーデータを入力することでモデル構造を確定させているというわけです。
また変換時のinput_namesやoutput_namesは推論時に利用します。入力・出力が複数ある場合、どこに何を入力すればよいか・どうやって取り出せばよいか分からなくなります。そういった問題を避けるために、各入力に対してラベルを振る作業をしています。
変換が終わればONNXを用いて推論をしましょう。ONNXを利用して推論を実行するにはONNXRuntimeが必要になります。早速インストールしましょう！
pip install onnxruntime
インストールが終わればいよいよ推論速度検証です。 以下のコードを実行しましょう！
from copy import deepcopy from time import time import cv2 import numpy as np import onnxruntime from utils.general import non_max_suppression def plot_bbox(frame, bboxes, time): for bbox in bboxes: bbox = bbox.astype(int) frame = cv2.rectangle(frame, bbox[0:2], bbox[2:4], (0, 255, 0), 3) cv2.putText(frame, text=f"time {time*1000:6.3f}", org=(0, 40), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1.0, color=(0, 255, 0)) return frame def cxcywh2xyxy(bboxes): bboxes_like = np.zeros_like(bboxes) bboxes_like[:, 0] = bboxes[:, 0] - bboxes[:, 2]/2 bboxes_like[:, 1] = bboxes[:, 1] - bboxes[:, 3]/2 bboxes_like[:, 2] = bboxes[:, 0] + bboxes[:, 2]/2 bboxes_like[:, 3] = bboxes[:, 1] + bboxes[:, 3]/2 return bboxes_like def main(): # カメラ接続 cap = cv2.VideoCapture(0) print("cam ok") # モデル読み込み。先ほど変換したonnxを利用 session = onnxruntime.InferenceSession("yolov5n6_onnx.onnx") print("model ok") # time list 30 frameの平均値を測る為 time_list = [0 for _ in range(30)] pointer = 0 while True: # 動画の1フレームを取得 ret, frame = cap.read() # 計測開始 start = time() # YOLOv5による推論 frame = cv2.resize(frame, (640, 512)) origin = deepcopy(frame) frame = np.float32(frame)[:, :, ::-1][np.newaxis, :, :, :]/255 frame = frame.transpose(0, 3, 1, 2) output = session.run(["output"], {"input": frame})[0][0, : ,:5] output = output[output[:, 4] > 0.8] #output = output[output[:, 6] ] #output = non_max_suppression(output)[0] #print(output) # 30 fraemの平均を出すために保存 time_list[pointer % 30] = time() - start pointer += 1 output = cxcywh2xyxy(output) #print(output) frame = plot_bbox(origin, output, np.mean(time_list)) cv2.imshow("test", frame) cv2.waitKey(1) if __name__ == "__main__": main()
ONNXのモデル読み込みには
session = onnxruntime.InferenceSession("yolov5n6_onnx.onnx")
を利用します。また推論では
session.run(["output"], {"input": frame})
とします。ここでONNX変換時に利用したラベルが生きていますね！
ONNXによる推論コードを書く際は、機械学習ライブラリに依存しないような書き方を意識しましょう。
その理由として以下の2点があります。
- モデルをONNXに変換したにも関わらず前処理でPyTorchを利用すると、結局PyTorchをインストールする羽目になり、複雑な環境構築から抜け出せなくなる
- 機械学習ライブラリからNumpy、Numpyから機械学習ライブラリへの変換に時間を要すため
では実行してみましょう。結果は・・・。
ここでダメ押しのOpenVINOを利用しましょう！
OpenVINOとは？
OpenVINOとはIntelが開発する推論ライブラリです。OpenVINOに変換することで、Intelプロセッサに最適化されたモデルに変換されるため、ほとんどのPCで高速化が見込まれます。ONNXRuntimeの推論速度で満足できなかったエッジAIエンジニアが行き着く先の一つです。
ちなみにそのほかの選択肢としてTensorRTを利用して高速化をする方法もあります。こちらはまた後日・・・。
ONNX vs OpenVINO
テレビの描画スピードである夢の33msまであと少しです。ここでさらなる高速化に向けてOpenVINOに手を染めましょう。
「お前もエッジAIエンジニアにならないか？」
OpenVINOの速度検証
では実際にOpenVINOを利用しましょう。まずはOpenVINOが利用できる環境を整えます。
pip install openvino-dev pip install openvino
続いてモデルの変換を行います。OpenVINOへの変換はたった１コマンドで実現できます。
mo --input_model yolov5n6_onnx.onnx
実行が完了すればいよいよ実験です。
#平均16FPS程度で動作 from copy import deepcopy from time import time import cv2 import numpy as np from openvino.inference_engine import IECore def plot_bbox(frame, bboxes, time): for bbox in bboxes: bbox = bbox.astype(int) frame = cv2.rectangle(frame, bbox[0:2], bbox[2:4], (0, 255, 0), 3) cv2.putText(frame, text=f"time {time*1000:6.3f}", org=(0, 40), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1.0, color=(0, 255, 0)) return frame def cxcywh2xyxy(bboxes): bboxes_like = np.zeros_like(bboxes) bboxes_like[:, 0] = bboxes[:, 0] - bboxes[:, 2]/2 bboxes_like[:, 1] = bboxes[:, 1] - bboxes[:, 3]/2 bboxes_like[:, 2] = bboxes[:, 0] + bboxes[:, 2]/2 bboxes_like[:, 3] = bboxes[:, 1] + bboxes[:, 3]/2 return bboxes_like def main(): # カメラ接続 cap = cv2.VideoCapture(1) print("cam ok") # モデル読み込み。先ほど変換したOpenVINOを利用 ie = IECore() model_path = 'yolov5n6_onnx.xml' weight_path = 'yolov5n6_onnx.bin' net_openvino = ie.read_network(model=model_path, weights=weight_path) model = ie.load_network(network=net_openvino, device_name='CPU', num_requests=1) print("model ok") # time list 30 frameの平均値を測る為 time_list = [0 for _ in range(30)] pointer = 0 while True: # 動画の1フレームを取得 ret, frame = cap.read() # 計測開始 start = time() # YOLOv5による推論 frame = cv2.resize(frame, (640, 512)) origin = deepcopy(frame) frame = np.float32(frame)[:, :, ::-1][np.newaxis, :, :, :]/255 frame = frame.transpose(0, 3, 1, 2) output = model.infer(inputs={'input': frame})['output'][0, : ,:5] output = output[output[:, 4] > 0.8] #output = output[output[:, 6] ] #output = non_max_suppression(output)[0] #print(output) # 30 fraemの平均を出すために保存 time_list[pointer % 30] = time() - start pointer += 1 output = cxcywh2xyxy(output) #print(output) frame = plot_bbox(origin, output, np.mean(time_list)) cv2.imshow("test", frame) cv2.waitKey(1) if __name__ == "__main__": main()
モデルの読み込み部分は
ie = IECore() model_path = 'yolov5n6_onnx.xml' weight_path = 'yolov5n6_onnx.bin' net_openvino = ie.read_network(model=model_path, weights=weight_path) model = ie.load_network(network=net_openvino, device_name='CPU', num_requests=1)
で構成されています。ONNXと比べると少し煩雑ですが、高速化できるのであれば何てことないですね！
推論は
model.infer(inputs={'input': frame})['output']
で行います。ONNXへの変換で付けたラベルがここでも利用されています。ラベル名は慎重に・・・。
では実行してみましょう。結果は・・・。
38ms！！！来ました！我々の勝利です！ほぼリアルタイム推論が出来るようになりました！
これでエッジに組み込んでいろんなことに使うことがことができそうです！
まとめ
今回はONNX・OpenVINOを用いてYOLOv5を高速推論可能なモデルへと変化させました！
| PyTorch | ONNX | OpenVINO | |
|---|---|---|---|
| 処理時間 | 150ms | 50ms | 38ms |
結果は上の表のようになり、OpenVINOが最も高速に推論可能であることが分かりました。
CPU上でリアルタイム推論が出来ることにより、様々な場面での活用が見えてくると思いますのでぜひお試しください！
