はじめに
個人で、スマホのカメラで麻雀牌を読み取るアプリを作っています。認識に使うモデルの候補は何十本もあるので、どれが速いかを本番の iPhone で自動計測する仕組みを作っていたのですが、これがよく途中で落ちました。落ちたときは「どうせ PC で測った速さと同じ傾向でしょ」と、手元の PC の数字で済ませていました。
ところが、同じモデルを PC と iPhone で測り比べてみると、モデル同士の差のつき方がまるで違いました。PC では1.8倍遅かったモデルが、iPhone では1割遅いだけだったのです。しかも PC のほうも、ONNX Runtime の最適化の設定を1つ変えるだけで、差のつき方が大きく変わりました。
PC の数字を信じていた私は、そのせいでモデル選びの幅を自分で狭めていました。この記事では、同じ ONNX モデルなのに、どうして環境や設定で傾向が変わるのかを追いかけた過程と、それで何を見落としていたのかを書きます。
目次
1. どうして PC で測っていたのか
本番では、モデルを iPhone の Safari 上で動かしています。ONNX Runtime Web という推論ライブラリを使って、ブラウザの中で WebAssembly として1スレッドで実行する構成です。
推論時間の計測は、学習が終わるたびに iPhone 上で自動で走る仕組みです。落ちる原因を毎回追いかけていると探索が止まってしまうので、学習に使っている PC の CPU でも同じ ONNX Runtime で並行して測っておき、iPhone の数字が取れなかったときはそちらで判断していました。
PC 側の計測条件
| 項目 | 条件 |
|---|---|
| CPU | Intel Core i5-11400(x86_64) |
| ONNX Runtime | CPUExecutionProvider |
| スレッド | intra / inter-op とも 1 |
| execution mode | ORT_SEQUENTIAL |
| graph optimization | ORT_ENABLE_ALL |
2. 「16の倍数を外れると遅い」と結論した実験
問題の実験では、チャネル数の違う CNN を何本か学習させて、精度と PC での推論時間を並べました。
| 各層のチャネル数 | 計算量(MAC) | PC での推論時間(p50) |
|---|---|---|
| 16/32/64/96 | 16.262M | 0.302 ms |
| 17/34/68/102 | 18.135M(+11.5%) | 0.541 ms(+79%) |
| 19/38/76/114 | 22.186M(+36.4%) | 0.605 ms(+100%) |
推論時間は 17/34/68/102 の時点で跳ね上がり、その先の 19/38/76/114 ではなだらかにしか増えていません。計算量(MAC、積和演算の回数)のほうは、どちらも素直に増えているだけなので、16の倍数のところにだけ崖があるように見えました。
19/38/76/114 は精度が少し良かったものの、推論時間が2倍になってまで選ぶほどではないと判断して、見送りました。
崖の理由は、生成AIに聞いて済ませました。ONNX Runtime の行列積は16個単位で計算していて、iPhone のブラウザで動く WebAssembly 版も同じ実装を使っている、という説明です。筋が通って聞こえたので、ソースも読まず、iPhone で測り直すこともしませんでした。
その後の実験でも、チャネル数は 16/32/64/96 に固定したまま、カーネルサイズや層の組み方など、別のところで性能を上げる方向に進めました。
3. 原因を探す
この記事を書くにあたって、原因を一から調べ直しました。
3.1 AIの説明は外れていた
まず、あのとき生成AIから聞いた説明が本当なのかを、ONNX Runtime のソースで確かめました。
ONNX Runtime の CPU 版は、計算の中身を MLAS というライブラリで行っています。MLAS の行列積(SGEMM)が、片方の行列を16列ずつまとめて詰め、余りを0で埋めているのは本当でした。AIの説明は、この「16列ずつ」がチャネル数の崖を作っている、というものです。
ところが Conv の処理を追っていくと、MLAS は入力を行列に並べ替えてから行列積を呼んでいて、16列ずつ詰められるのは出力ピクセルの方向でした。チャネルの方向ではないので、チャネルを16から17にしても詰め方は変わりません。崖の説明としては外れです。
3.2 当時のモデルで、まず再現させる
ソースを読むだけで原因を当てるのは難しそうだったので、実際に動かして確かめることにしました。当時測った ONNX ファイルは実験管理ツール(ClearML)に残っていたので、それを取り出して同じ PC で測り直すと、当時とほぼ同じ値が出ました(16幅 0.309ms、17幅 0.554ms。ONNX Runtime 1.28.0)。崖は再現したので、この2つのモデルの中で何が起きているかを見ていきます。
3.3 実行前のモデルを覗く
ONNX Runtime は、読み込んだモデルをそのまま実行するのではなく、速く動くように組み替えてから実行します。組み替えた後のモデルはファイルに書き出せるので、16幅と17幅で中身を比べてみました。
| ノード | 16/32/64/96 | 17/34/68/102 |
|---|---|---|
| NCHWc の Conv | 5 | 2 |
| 普通の Conv | 0 | 3 |
| NCHWc の MaxPool / GlobalAveragePool | 全部 | — |
| ReorderInput | 0 | 1 |
| ReorderOutput | 1 | 2 |
違いは「NCHWc」という最適化でした。チャネルを16個などのまとまりに束ねて、CPU の SIMD 命令でまとめて計算しやすくするもので、x86 などの CPU でだけ使われます。ただし変換できるかどうかにはチャネル数の条件があり1、16/32/64/96 はモデル全体が NCHWc で動くのに対して、17/34/68/102 は5層のうち3層が条件から外れて普通の Conv のまま残り、その間でデータの並びを行ったり来たりさせる処理(Reorder)が挟まっていました。
3.4 設定を1つ変えるだけで、崖が消える
NCHWc が本当に原因なのかは、NCHWc を使わずに動かしてみれば分かります。最適化のレベルを1段下げると(ORT_ENABLE_EXTENDED)、NCHWc への変換は行われません。
| 各層のチャネル数 | ENABLE_ALL | ENABLE_EXTENDED |
|---|---|---|
| 16/32/64/96 | 0.309 ms | 0.463 ms |
| 17/34/68/102 | 0.554 ms | 0.493 ms |
| 差 | +79% | +6.5% |
同じ PC、同じモデルなのに、設定を1つ変えただけで、差は +79% から +6.5% まで縮みました。崖のほとんどは、NCHWc が効くかどうかで説明できます。
もう1つ面白いのは、17幅が遅くなったというより、16幅のほうが NCHWc で約1.5倍速くなっていたことです(0.463 → 0.309ms)。17幅は逆に、NCHWc を止めたほうが速くなっています(0.554 → 0.493ms)。中途半端に変換されるくらいなら、変換しないほうがましだったわけです。
4. iPhone で測ってみる
NCHWc は x86 などの CPU 向けの最適化で、ONNX Runtime のソースを見ると、WebAssembly では有効になりません。だとすると、本番の iPhone ではこの崖は出ないはずです。当時の ONNX ファイルを、今度は iPhone で動かしてみました。
iPhone 側の計測条件
| 項目 | 条件 |
|---|---|
| 端末 | iPhone 13 |
| ブラウザ | Safari |
| ONNX Runtime Web | 1.27.0 |
| Execution Provider | wasm-simd |
| スレッド | 1 |
| graph optimization | all |
| batch | 1 |
4.1 測る順番で結果が変わる
最初はモデルを1本ずつ順番に測ったのですが、後に測ったモデルほど遅く出てしまいました。測っているうちに端末が温まって、だんだん遅くなっていくためです。
そこで、1つのページに3つのモデルを同時に読み込み、100回の推論を1ブロックとして、次の順番を20周回しました。
16 → 17 → 19 → 19 → 17 → 16
左右対称に挟んでいるので、各モデル40ブロックずつ、温度の変化を同じように受けることになります。それでも端末は少しずつ温まっていくので2、ここで見るのは絶対値ではなく、モデル同士の比です。
4.2 iPhone には崖がなかった
| 各層のチャネル数 | 計算量(MAC) | iPhone(p50) |
|---|---|---|
| 16/32/64/96 | 16.262M | 1.18 ms |
| 17/34/68/102 | +11.5% | +10.2% |
| 19/38/76/114 | +36.4% | +29.7% |
iPhone では、17幅が +10.2%、19幅が +29.7% で、どちらも計算量の増え方とほぼ同じでした3。
ここまでの結果を、16幅に対して17幅がどれだけ遅いかで並べると、同じモデルでも環境と設定によって傾向がまるで違うことが分かります。
| 環境と設定 | 17幅の遅さ(16幅比) | 傾向 |
|---|---|---|
| PC(x86)、最適化 ENABLE_ALL | +79% | 16の倍数のところに崖がある |
| PC(x86)、最適化 ENABLE_EXTENDED | +6.5% | ほぼ計算量どおり |
| iPhone(WASM) | +10.2% | ほぼ計算量どおり |
PC で見えていた崖は、PC の、しかも特定の最適化設定の中にしかなかったことになります。
5. 何を見落としていたか
16/32/64/96 は iPhone でも17幅や19幅より速いので、16の倍数に揃えたこと自体は、結果的に間違いではありませんでした。間違っていたのは理由のほうで、16/32/64/96 は単にチャネルが少ないから速く、PC ではそこに NCHWc のおまけが乗っていただけでした。
この違いがはっきり出るのが、19/38/76/114 です。19幅の精度は、歪みや切り出し位置のずれを加えた条件ごとの正解率で見ると、平均 0.902、一番悪い条件で 0.733 と、16幅(0.889、0.722)より少し良い結果でした。
| 19幅の推論時間(16幅比) | |
|---|---|
| 当時見ていた PC | 約2倍 |
| 本番の iPhone | 約1.3倍 |
当時は「精度が少し上がる代わりに2倍遅い」モデルに見えていましたが、本番の iPhone で見れば「精度は同等以上で3割遅い」モデルです。候補として検討する価値は、当時考えていたよりありました。もっとも、精度は学習1回分の結果で差も 0.01〜0.02 なので、良いモデルを逃したとまでは言えません。
ちなみに今は、モデルを選ぶときに iPhone での推論時間を毎回見て決めています。
おわりに
同じ ONNX モデルでも、動かす環境と最適化の設定によって、速さの傾向は大きく変わりました。PC の CPU では、ONNX Runtime が x86 向けの最適化をかけられるかどうかで崖が出たり消えたりし、iPhone のブラウザではその最適化がそもそもなく、差はほぼ計算量どおりでした。どれかの環境の傾向が正しいわけではなく、環境ごとに別の傾向が出る、というだけです。
なので、代わりの環境で測った数字を使うなら、本番と同じ傾向が出るかを一度は確かめておく必要がありました。私の場合は、生成AIのもっともらしい説明を聞いて、その確認を飛ばしてしまいました。ONNX Runtime や WebAssembly のソースを読み解くのは大変ですが、同じモデルを本番の端末で測るだけなら誰にでもできますし、実際、iPhone で一度測っただけで崖がないことは分かりました。
今回確かめたのは、小さな CNN 1種類と、PC 1台、iPhone 1台だけです。ほかのモデルや端末で同じことが起きるかは分かりません。ただ、少なくとも「PC で測っておけば、スマホでも同じ傾向だろう」とは言えなさそうです。製品の判断に直結する数字ほど、本番の環境で測って確かめる。今回、身をもって学んだのはそれでした。
-
主な条件は、入力チャネルがブロック(この PC では16)以上の Conv では入力チャネルが4の倍数であること、Pool ではチャネルがブロックの倍数であることです。出力チャネルはブロックの倍数まで0で埋められます。17幅で外れたのは、入力チャネルが 17、34、102 の層でした。変換された層も0埋めで大きくなり、1→17 の層は内部の重みが 32ch、68→102 の層は 68→80、102→112 相当になっていました。 ↩
-
室温27℃、端末は扇風機の上に置いて測りましたが、最初と最後で16幅は 1.05〜1.10ms から 1.18〜1.20ms に、19幅は 1.38〜1.43ms から 1.53〜1.55ms に上がっていました。 ↩
-
17幅は、別の日に単独で測ったときも +10.7% でした。 ↩