ONNX推論クレートの比較と実装奮闘記

Rust.Tokyo 2024で発表した資料です。
https://rust.tokyo/lineup/13

Transcript

1. ONNX推論クレートの比較 と実装奮闘記 2024-11-30 Satoshi Yoshikawa 1 アプリケーションへの組み込み 機械学習・AIのモデルフォーマット
2. 本日のトピック • 自己紹介 • 今回のお話の背景 • ONNXと推論エンジンについて • 自身の取り組み（実装奮闘記の部分） 2
3. 本日のトピック • 自己紹介 • 今回のお話の背景 • ONNXと推論エンジンについて • 自身の取り組み（実装奮闘記の部分） 3
4. 自己紹介 • X (ex. Twitter)/GitHub: @emergent • Fairy Devices株式会社 プロダクト開発部 部長 • Publication ◦ 『実践Rustプログラミング入門』共著 ◦ WEB+DB PRESS No.129 インタビュー ◦ WEB+DB PRESS No.131 特集「Rust入門」 ◦ 『RustによるWebアプリケーション開発 設 計からリリース・運用まで』共著 Satoshi Yoshikawa / 吉川哲史 new! 4 自己紹介
5. 著書の宣伝 蔵書管理アプリケーションをテーマに、Rust でWebアプリケーションを開発するステップ を本にしました よかったらお手にとってみてください 5 RustによるWebアプリケーション開発 設計からリリース・運用まで 豊田 優貴, 松本 健太郎, 吉川 哲史 https://www.amazon.co.jp/dp/4065369576 自己紹介
6. 本日のトピック • 自己紹介 • 今回のお話の背景 • ONNXと推論エンジンについて • 自身の取り組み（実装奮闘記の部分） 6
7. 本日のお話の背景 • 当社の研究開発部門から音声処理の機械学習モデル （ONNX形式）が提供された • せっかくなのでRustで推論処理を実装したい • あわよくば自社デバイス上で動かせるようにしたい 7 本日のお話の背景
8. 本日のお話の背景 • 当社の研究開発部門から音声処理の機械学習モデル （ONNX形式）が提供された • せっかくなのでRustで推論処理を実装したい • あわよくば自社デバイス上で動かせるようにしたい 8 しました、という話です 本日のお話の背景
9. 本日のトピック • 自己紹介 • 今回のお話の背景 • ONNXと推論エンジンについて • 自身の取り組み（実装奮闘記の部分） 9
10. 本日のトピック • 自己紹介 • 今回のお話の背景 • ONNXと推論エンジンについて ◦ ONNXとは ◦ ONNXの推論エンジン ◦ ONNX形式のモデルを扱えるクレート • 自身の取り組み（実装奮闘記の部分） 10
11. ONNX（Open Neural Network Exchange） • https://onnx.ai/ • 機械学習モデルを表現するためのフォーマット • フレームワークの相互運用性・最適化の共有を目的に開発 されている ◦ 学習時のフレームワークに依存せず推論できる ◦ 推論対象のフォーマットの統一によりパフォーマンス調整の対象 を限定できる 機械学習フレームワーク （TensorFlow、PyTorch など） export ONNX モデル 推論エンジン import 11 ONNXとは
12. ONNXフォーマットの内容 • 定義されたオペレーター（演算子）をノードとして、入 出力でグラフ状に接続 • モデルのグラフはProtocol Buffer（protbuf）でシリアラ イズされて保存される 12 ONNXとは
13. オペレーター • Inputs（入力） • Initializer（初期化子） • Attributes（属性） • Outputs（出力） 13 右図はNetronで可視化したMNISTモデル Netron https://github.com/lutzroeder/netron ONNXとは
14. Opsetバージョン • ONNXモデルがどのバージョンのオペレーターで構成され ているかを示す情報 • 同じオペレーターでもOpsetバージョンによって仕様が異 なる場合がある ◦ Opsetバージョンごとの差分はONNX Operatorsにて確認できる 14 更新があったバージョン バージョンごとの差分 ONNXとは
15. 本日のトピック • 自己紹介 • 今回のお話の背景 • ONNXと推論エンジンについて ◦ ONNXとは ◦ ONNXの推論エンジン ◦ ONNX形式のモデルを扱えるクレート • 自身の取り組み（実装奮闘記の部分） 15
16. ONNX形式のモデルを推論できるエンジン • ONNX Runtime • HWと紐づいた推論エンジン ◦ OpenVINO（Intel）、TensorRT（NVIDIA） • その他の推論エンジン ◦ →Rust製の推論エンジンを次項で紹介 16 ONNXの推論エンジン
17. ONNX Runtime • https://onnxruntime.ai/ • ONNXモデルで推論を実行するエンジン • Microsoftがオープンソースで開発・維持している • C++で実装 • Execution Provider ◦ HWアクセラレーションを使う仕組み 17 ONNXの推論エンジン
18. 本日のトピック • 自己紹介 • 今回のお話の背景 • ONNXと推論エンジンについて ◦ ONNXとは ◦ ONNXの推論エンジン ◦ ONNX形式のモデルを扱えるクレート • 自身の取り組み（実装奮闘記の部分） 18
19. ONNX形式モデルを推論できるクレート 19 クレート名 構成 扱う内容 説明 ort バインディング 推論 ONNX Runtimeをラップしたもの tract Pure Rust 推論 Sonos社が開発している推論エンジン wonnx Pure Rust 推論 ready for the webを謳う推論エンジン burn Pure Rust 学習・推論 Rust製の機械学習フレームワーク candle Pure Rust 学習・推論 Rust製の機械学習フレームワーク kyanite Pure Rust 推論 個人開発の推論エンジン rten Pure Rust 推論 個人開発の推論エンジン ONNX形式のモデルを扱えるクレート
20. ONNX形式モデルを推論できるクレート • MNISTモデルと手書き画像を用意して 各クレートで推論させてみた ◦ リポジトリ ：https://github.com/emergent/compare-onnx-runtimes 実行環境 　MacBook Pro - Apple M1 Max 　macOS Sonoma 14.7 20 ONNX形式のモデルを扱えるクレート
21. 用意した画像 21 どこからどう見ても人の目には 「３」 に見える手書き画像 MNISTの入力仕様に合わせて 28x28のサイズのVec<f32> に調整した ONNX形式のモデルを扱えるクレート
22. 推論のやり方（共通） 1. 入力の画像データからテンソルを作る 2. モデルを読み込み推論セッションを初期化 3. テンソルを入力して推論を実行 4. 出力から必要な情報を取り出す 22 ONNX形式のモデルを扱えるクレート
23. 実装例（ortを使った実装） 23 1. 入力の画像データから テンソルを作る 2. モデルを読み込み 推論セッションを初期化 3. テンソルを入力して 推論実行 4. 出力から必要な 　 情報を取り出す リポジトリ：https://github.com/emergent/compare-onnx-runtimes ONNX形式のモデルを扱えるクレート
24. MNISTで推論させた結果（初期化時間は含まない） 24 クレート名 推論処理時間（ms） 結果の補足 ort 0.35 はやい tract 0.90 はやい wonnx 3.44 ちょっとおそい burn 対象外 正しい推論結果を出せなかった candle 実行不可 auto_pad: SAME_UPPERに非対応 kyanite 1.51 ふつう？ rten 0.66 はやい（モデルを事前変換しているため？） ONNX形式のモデルを扱えるクレート
25. 自社のモデル実装に使うクレート選定 • 実装の目的 ◦ 動作対象は自社のウェアラブルデバイス〜PCまたはサーバー ◦ GPU処理は不要 • →ort か tract に絞って実装にチャレンジする ◦ ort → ONNX RuntimeのExecution Providerを使える（かも） ◦ tract → 自社デバイスで使っているCPU向けの 　　　　 最適化処理が実装されている（ようだ） 25 ONNX形式のモデルを扱えるクレート
26. ウェアラブルデバイス THINKLET 26 https://mimi.fairydevices.jp/technology/device/thinklet/ Androidアプリを作れば動く ONNX形式のモデルを扱えるクレート
27. 本日のトピック • 自己紹介 • 今回のお話の背景 • ONNXと推論エンジンについて • 自身の取り組み（実装奮闘記の部分） 27
28. 本日のトピック • 自己紹介 • 今回のお話の背景 • ONNXと推論エンジンについて • 自身の取り組み（実装奮闘記の部分） ◦ 目的のおさらい ◦ tractでの実装 ◦ ortでの実装 ◦ まとめ 28
29. 今回の実装ゴール 1. 音声処理のためのONNXモデルを 2. Androidアプリに組み込み 3. リアルタイム動作可能なレベルに持っていく 29 目的のおさらい
30. 今回の実装ゴール 1. 音声処理のためのONNXモデルを 2. Androidアプリに組み込み 3. リアルタイム動作可能なレベルに持っていく 30 とは？ 目的のおさらい
31. 今回の実装ゴール 1. 音声処理のためのONNXモデルを 2. Androidアプリに組み込み 3. リアルタイム動作可能なレベルに持っていく 31 とは？ RTFの値が 1.0 より十分小さい RTF：Real Time Factor 目的のおさらい
32. RTF（Real Time Factor）とは 「1 秒の入力データが与えられたとき、処理が完了するため に、その何倍の時間が掛かるか？」という指標 32 RTF ＝ 処理時間（秒） 音声データ長（秒） 目的のおさらい
33. RTF（Real Time Factor）とは 33 詳細はこちら https://mimi.fairydevices.jp/technology/cloud/asr/ 音声処理 音声処理 目的のおさらい
34. 本日のトピック • 自己紹介 • 今回のお話の背景 • ONNXと推論エンジンについて • 自身の取り組み（実装奮闘記の部分） ◦ 目的のおさらい ◦ tractでの実装 ◦ ortでの実装 ◦ まとめ 34
35. 最初はtractで実装を試みる • Pure RustなのでAndroidアプリ組み込み時の面倒事が少 なそうであるため 35 tractでの実装
36. 処理フロー 初期化 音声処理 （一定サイズごと） 処理済み音声出力 ONNXモデル読み込み 音声データに前処理を行い 推論処理にかける ファイルに出力 音声再生 など 36 tractでの実装 入力音声が 終了するまで 繰り返す
37. PC上で動作をさせてみる • PC上でのRTFは0.1程度 ◦ 簡易的なCLIを作って自身のPC上で音声ファイルを処理 ◦ ラズパイ4だと0.4、ラズパイ5だと0.15ぐらい • これはTHINKLETでも軽快に動作させられるのでは？ （淡い期待） 37 tractでの実装
38. Androidアプリへの組み込み 1. JNIインタフェースを実装（jni-rs） 2. mozilla/rust-android-gradle を 使ってAndroidプロジェクトに組み込む 3. Kotlin側のクラスを実装 4. 簡単なUIを実装してVysorやscrcpy経由 で画面を見られるようにする 38 tractでの実装
39. アプリの初期実装シーケンス 39 アプリ（Kotlin） JNI Rust 推論セッション の初期化 インスタンス生成 音声データ入力 （チャンク単位） 推論処理の実行 推論結果の取得 元音声を適切な単位 （チャンク）に分割 し繰り返し実行する 最初はすべて同期的に実行する形で実装 tractでの実装
40. THINKLET上で動作させる • RTFは0.95程度（初期化処理は除く） ◦ リアルタイム動作させるにはギリギリ ◦ もう少し軽快に動作させたい • topで処理を覗いてみると1コアしか使っていない ◦ PCよりしょぼい環境と言えどマルチコアなので有効活用したい 40 tractでの実装
41. 力技でマルチスレッド処理にする スレッドを複数作成し、各スレッドそれぞれで 独立に初期化・推論処理を行うようにした 推論スレッド 推論スレッド … 出力スレッド （コールバック） 入力スレッド 41 tractでの実装 矢印部分は mpsc::channel を使用しデータを転送
42. 力技マルチスレッドの結果 • 推論スレッド2でRTF 0.6程度、3で0.5を切れた ◦ スレッド数を増やすほど初期化時間がかかる（わかる） • 音声データの投げ方に工夫は必要だったが効果はあった 42 参考： CPUは少し前の世代の 8コアのもの tractでの実装
43. tractのCPU最適化処理を適用 • 使われてない（自動判別されてない）ことに気づいた • 環境変数設定で該当のCPUを強制設定→有効化できた • スレッド数を増やすとRTFの改善効果そのものは薄れる 43 スレッド数1で 約15％の改善 tractでの実装
44. 本日のトピック • 自己紹介 • 今回のお話の背景 • ONNXと推論エンジンについて • 自身の取り組み（実装奮闘記の部分） ◦ 目的のおさらい ◦ tractでの実装 ◦ ortでの実装 ◦ まとめ 44
45. ortを使ってみる • PCでは軽快に動作した（tractでも軽かったが） • Androidアプリに組み込んでどうなるか 45 ortでの実装
46. • ONNX Runtime ◦ ビルド手順：https://onnxruntime.ai/docs/build/android.html • ort ◦ 上記とリンクさせるため、ortもクローンしてビルド ▪ 環境変数 ORT_LIB_LOCATION, ORT_LIB_PROFILEを指定 ◦ Cargo.tomlからはpath指定してビルド ▪ • アプリには libc++_shared.so も同梱が必要 ◦ アプリ内にjniLibs/arm64-v8aディレクトリを掘って保存 ONNX Runtime / ort を自分でビルドする 46 ort = { version = "2.0.0-rc.8", path = "../ort" } ortでの実装
47. ort版を動かしてみた • tractに比べてモデル読み込みがかなり高速（約4倍） • 同期処理での推論は少し速い程度 • 一方で、 ◦ NNAPI（※）有効化も試したが大きな改善は見られず ◦ 力技マルチスレッド処理も（同期シーケンスよりは改善したが） スレッド数を増やすことによる改善は見られなかった 47 ortでの実装 ※NNAPI（Neural Networks API）：Android デバイス上で演算負荷の高い 　機械学習処理を実行するために設計された Android C API 　ONNX RuntimeのExecution Providerとして使うことができる
48. 上手くいかなかったもの • tract, ortいずれも、クレートがもつ並列化機能を使って も速度向上できなかった ◦ tract ▪ マルチスレッド機能で指定するスレッド数を増やすほど遅くなった ◦ ort ▪ intra_op_threads, inter_op_threads というパラメーターを 指定できるが、いずれも「指定しない」が最も速かった ▪ ortは「指定しない」で実行しても複数コアで実行できて いるようではあった 48 モデル内に並列化で効率的に実行できるオペレーター が（ほとんど）なかったから？ ortでの実装
49. 本日のトピック • 自己紹介 • 今回のお話の背景 • ONNXと推論エンジンについて • 自身の取り組み（実装奮闘記の部分） ◦ 目的のおさらい ◦ tractでの実装 ◦ ortでの実装 ◦ まとめ 49
50. 試したチューニングまとめ 50 クレート やったこと 効果 説明 tract マルチスレッド機能 ✕ 遅くなった 力技マルチスレッド ◯ 効果あり CPU最適化の有効化 △ スレッド数が少なければ効果あり ort intra_op_threads ✕ 遅くなった inter_op_threads ✕ 遅くなった NNAPI ✕ 改善なし 力技マルチスレッド △ 同期処理よりはましな程度 ※ortの場合は「いらんことしない」が正解かも？ まとめ
51. とりあえずの結論 • どちらのクレートでもAndroidアプリに組み込んで動作さ せ、速度計測はできた • いくつかの方法でRTFの数値改善ができたが、tractの方 がコントロールしやすかった • 最終的な実装は製品化に向けてもう少し調査したい 51 まとめ
52. 今後やってみたいこと • ONNX Runtimeの並列処理をもっと理解する • モデルそのものの改善（軽量化）と評価 • 製品化のためのライブラリ化 • 他言語バインディングの実装 • ONNXにエクスポートしない場合での推論との速度比較 52 まとめ
53. 全体のまとめ • ONNXについて調べた • ONNXでの推論処理クレートを試した • 実際に自社デバイス向けアプリに実装して動かした 53 まとめ
54. ご清聴ありがとうございました 54 Zzz…