軸がひとつ固まるとバッチが通らない

一言でいうと

ONNXの軸は、整数で固定されているか、シンボル名で開いているか、そもそもないかの3つのうちのどれかです。バッチを大きくできない事故の大半は、3つ目ではなく1つ目から来ます。しかも、エラーなしで通り過ぎる箇所から始まります。

なぜ必要なのか

量子化したモデルがどれだけ速くなったかを測るには、バッチを大きくしてみる必要があります。32個を一度に入れようとしたら、ランタイムに拒否されます。モデルは確かに「動的バッチ」でエクスポートしたと聞いていて、セッションもきちんと開きました。ところが、値を入れた瞬間に止まります。

ここでよくある対応は、ランタイムのオプションを探し回ることです。答えはファイルの中にあります。セッションが開くということは、グラフが成り立つという意味にすぎず、私たちが入れようとしているシェイプを受け付けるという意味ではありません。

軸は3つの状態を持つ

ONNX Conceptsが説明するテンソルのシェイプは、軸ごとに3つのうちのどれかです。

• 整数: dim_valueが書き込まれています。その軸は、その値でなければなりません。
• シンボル名: dim_paramに文字列が入っています。実行時に決まり、同じ名前は同じ値を意味します。
• 未指定: 両方ともありません。不明という意味で、推論もここで止まります。

2つ目の「同じ名前は同じ値」は、よく見落とされます。2つの入力がどちらも軸0にbatchと書かれている場合、それは単なる「両方とも動的」ではなく、2つの行数が互いに同じでなければならないという制約です。片方に5行、もう片方に3行を入れると拒否されます。

ただし、その拒否は親切ではありません。ランタイムは「シンボルbatchが5と3に分かれた」とは教えてくれません。グラフ最適化を経たあとのノード名を挙げて、そのノードの入力シェイプが合っていないと伝えます。原因と症状の間に最適化という1層が挟まっているので、初めて見ると原因が見えません。

선언            x [batch, 4]   bias [batch, 3]
넣은 것          x 5행          bias 3행
런타임의 말      (융합된 노드 이름) 의 입력 모양이 맞지 않는다
실제 원인        같은 심볼에 다른 값을 넣었다

推論が埋めるものと埋められないもの

onnx.shape_inferenceは、グラフをたどって中間テンソルのシェイプを埋め、value_infoに入れてくれます。入力が[batch, 4]で、重みが4 x 6なら、中間テンソルは[batch, 6]と埋められます。シンボルはシンボルのまま伝播されます。

埋められない場合が2つあります。1つは、入力のシェイプが未指定のときです。わからないものから出発すれば、最後までわかりません。もう1つは、中間テンソルがそもそもないときです。ノードが1つだけで、その出力がそのままグラフの出力になる場合は、埋める対象がないので、value_infoが空のまま返ってきます。このとき「推論が失敗した」と読んではいけません。埋めるものがなかっただけです。

さらに危険な場合があります。推論が誤った値を自信たっぷりに埋めてしまうことです。Reshapeのターゲットシェイプが定数で埋め込まれていると、推論はその定数をそのまま信じます。入力を[batch, 4]と開いておいても、Reshapeが[3, 6]を指していれば、その後のすべてのテンソルは3で固定されます。エラーは出ません。バッチ軸は宣言にだけ残り、実際には死んでしまいます。

現場での姿

1つ目は、バッチ1でエクスポートしたモデルです。エクスポートツールに例の入力を1つ渡すと、そのシェイプがそのまま固定されます。動的軸を指定しないと[1, ...]が埋め込まれ、そのファイルはいつまでも一度に1件しか処理できません。スループットを測る実験が丸ごと無意味になります。

2つ目は、拒否がセッションを開くときではなく、値を入れるときに起きることです。セッションは問題なく開きます。そのため「モデルのロードはできたので、モデルの問題ではない」と結論づけて、見当違いの場所を掘ってしまいます。固定された軸のエラーメッセージは、どの入力の何番目の軸で何を期待し何を受け取ったのかを正確に書いてくれるので、その行だけ読めば済みます。

3つ目は、定数のシェイプが埋め込まれたReshapeです。これが最も静かな事故です。モデルを開くと、入力は[batch, 4]と問題なく開いていて、推論結果も辻褄が合っています。ただ、バッチ軸があるところから整数に変わっています。そのため「動的軸でエクスポートしたのに、なぜ動かないのか」という声が出続けます。直し方は単純で、バッチの位置を-1にすると、残りの軸から計算されて軸が生き返ります。

4つ目は、量子化は軸を固定しないものの、原因を隠すことです。量子化されたグラフは、ノード名が変わってQ/DQが入り込み、読みにくくなります。そのため、軸の問題は量子化の前に確認しておくほうがはるかに安上がりです。

実務で本当に大切なこと

• 入れてみてから話します。宣言だけを読まず、さまざまな行数で実際に入れてみた記録を残します。
• 同じシンボルは同じ値です。2つの入力が同じ名前を使っているなら、その制約をドキュメントに書きます。
• 推論の空白と確信を区別します。空のままのものと、誤って埋められたものは、まったく別の問題です。
• バッチの位置は-1にします。Reshapeのターゲットシェイプに、バッチの値を直接書きません。

次のラボですること

同じ重みで、バッチ軸がシンボルのバージョンと、1で固定されたバージョンを作って並べ、ツールaxes.pyをステップごとに育てます。2つのファイルに同じコードでさまざまな行数を入れて、どこで分かれるかを記録し、推論が埋めたものと埋められなかったものを分け、同じシンボルを使う2つの入力に異なる行数を入れてみます。最後に、定数のシェイプが埋め込まれたReshapeを見つけ出し、バッチの位置を元に戻して軸を生き返らせます。採点ツールは、毎回異なる軸名とシェイプ、異なる埋め込み行数で自分のファイルを作り、皆さんのツールを実際に動かして答えを照合します。