前回は、バッチサイズ1の入力のみを受け付けるCore MLモデルを生成し、実行できることを確認しました。 DNNの実行効率(1サンプル当たりの所要時間)は一般的にバッチサイズが大きいほど向上します。 そのため、DNNを効率よく実行するにはバッチサイズ2以上の入力を受け付けるようなCore MLモデルを生成したいです。 その手段として、(1)バッチサイズ256のような固定のバッチサイズ(>1)を受け入れるモデルを生成する、(2)実行時に任意のバッチサイズを受け入れるモデルを生成するという2通りが考えられます。 コンピュータ将棋の場合、固定のバッチサイズに満たないサンプル数でモデルを実行したい場合もある(探索開始時にルート局面だけを評価する場合等)ため、(2)が実現できると効率が良いのでこれを実現できるか試しました。 任意のバッチサイズの入力を受け取って処理できるモデルをここでは「バッチサイズ可変モデル」と呼ぶことにします。
結果の要約は以下のようになります。
- バッチサイズ1024を与えたときの処理時間が約1秒
- 1サンプル当たりの計算量は1GFLOPs。Neural Engineの処理速度はおおむね1TFLOPSと見積もれる。
- 将棋のゲーム木探索に用いることを考えると、1000NPSでありかなり低め。もっと計算量が小さいモデルが良い。
- わずかな条件差でCore ML内部で実行時エラーが起きやすく扱いづらい
- バッチサイズを切り替えるたびに余分な時間がかかるため非効率
バッチサイズ可変モデルの生成
バッチサイズ可変モデルは、coremltoolsでの変換時にオプションを与えることで実現できます（ライブラリのマニュアル）。ライブラリの機能上、バッチサイズだけでなく画像サイズなどを可変にする場合にも使用できます。
ライブラリに対して、モデルの入力のどこが変化する可能性があるか教える必要があります。その柔軟性に応じて区分があります。
- 入力しうる形状を列挙する方法（列挙タイプ）
- ある軸の長さの範囲を指定する方法（範囲タイプ）
列挙タイプでは、形状を可変にできる入力テンソルが1個という制約があります。しかしdlshogiのモデルではx1, x2という2つの入力があり、それぞれ[バッチサイズ, 62, 9, 9]、[バッチサイズ, 57, 9, 9]という形状のため両方を可変にしなければならずそのままでは実装できません。そのため、入力変数を1個にまとめて、モデル内部で分割してx1, x2として扱うようにしました。コードは以下のようになります。
class SingleInputModel(torch.nn.Module): def __init__(self, model): super().__init__() self.model = model def forward(self, x): return self.model(x[:, 0:62, :, :], x[:, 62:119, :, :])
列挙タイプの場合は、形状を表すタプルのリストをEnumeratedShapesクラスに与えます。
input_shape = ct.EnumeratedShapes(shapes=[(b, 119, 9, 9) for b in [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]]) ct_input = ct.TensorType(name='x', shape=input_shape) mlmodel = ct.convert(traced_model, inputs=[ct_input])
範囲タイプの場合は、形状を変化させたい軸にRangeDimクラスを用います。ここで、RangeDimの引数を空にすれば任意の数を受け取ることを表現できます。今回はそこまでの自由度は必要ないため十分な大きさとして1024を与えました。
input_shape = ct.Shape(shape=(ct.RangeDim(1, 1024), 119, 9, 9)) ct_input = ct.TensorType(name='x', shape=input_shape) mlmodel = ct.convert(traced_model, inputs=[ct_input])
列挙タイプのメリットは、指定された各形状にあわせた最適化がなされるとドキュメントに記載があります。
バグの回避
これで単純に変換できれば簡単だったのですが、coremltoolsのバグを踏んでしまうようでValueError: Cannot add const [is35, 2187]というエラーが発生しました。調査したところ、flattenオペレータの変換部分で失敗していることがわかりました。dlshogiモデル内でのflattenオペレータの使用状況を調べ、出力の形状を決め打ちで与えるオペレータで上書きしてやることにより、バグを回避することに成功しました。（簡単そうに書いていますが調査はかなりの知識を要します）
# dlshogiモデルに特化した内容であり汎用性はないので注意
from coremltools.converters.mil.frontend.torch.torch_op_registry import register_torch_op
from coremltools.converters.mil.frontend.torch.ops import _get_inputs
from coremltools.converters.mil.mil import Builder as mb
@register_torch_op(override=True)
def flatten(context, node):
inputs = _get_inputs(context, node)
x = inputs[0]
dims = [-1, 2187] # 決め打ち設定。reshapeStaticオペレータが出力されるが、長さ-1の次元は許容される
reshape = mb.reshape(x=x, shape=dims, name=node.name)
context.add(reshape)
以上のようなテクニックを組み合わせ、Core ML形式のモデルを得ることができました。colab notebook
実行
前の章で生成したモデルをSwift側で実行する方法ですが、バッチサイズ固定の場合と違いはありません。与えるMLMultiArrayに形状情報が含まれているので、バッチサイズ等の明示的な情報付与は不要です。変換の都合で入力変数がx1, x2の2つからxの1つに変わったので、そこは書き換えます。
guard let pred = try? model.prediction(x: sampleIO.x) else {
msg = "Error on prediction"
return
}
簡易ベンチマーク
実行速度を簡易的にベンチマークしました。あくまで1回の試行結果です。入力バッチサイズを変えると初回の実行が遅くなる場合がある（後述）のため、2回目以降の処理時間を示します。
範囲タイプのモデルの処理時間を示します。
| 環境 | バッチサイズ | 処理時間[s] |
|---|---|---|
| NE | 1 | 0.007 |
| NE | 8 | 0.025 |
| NE | 16 | 0.033 |
| NE | 64 | 0.086 |
| NE | 256 | 0.26 |
| NE | 1024 | 0.99 |
| CPU | 1 | 0.041 |
| CPU | 8 | 0.090 |
| CPU | 64 | 0.22 |
| CPU | 1024 | 3.15 |
処理時間は、model.predictの前後でDate()により得た時刻の差分で測定しています。「環境」は、NE=Neural Engine (computeUnits = .all)、GPU(computeUnits = .cpuAndGPU)、CPU(computeUnits = .cpuOnly)の3種類です。GPUを指定すると、以下のような実行時エラーが発生したため計測できませんでした。NE, CPUではともに正常動作しているため、モデルの誤りではなくCore ML内部の実装のバグではないかと思われます。
validateComputeFunctionArguments:745: failed assertion `Compute Function(TARR_elementwise_add_broadcast_f16_pack4): missing buffer binding at index 1 for p2[0].'
変換元となったdlshogiのモデルは10ブロック（畳み込み層20回）、192チャンネルのものです。計算量は1.099GFLOPsです（算定方法：畳み込み層と全結合層のみ。積和演算を2回の演算とみなす。）。どのようなバッチサイズでも、Neural EngineのほうがCPUより高速です。Neural Engine、バッチサイズ1024のときが最も高速で、1.137TFLOPSの演算能力が得られることがわかりました。将棋AIとして探索と組み合わせることを考えるとバッチサイズは16程度が現実的で、この場合で533GFLOPSとなります。モバイル端末としては十分な能力なのですが、将棋AIとしては1000NPS程度しか出ないというのは探索不足で弱いと予想されます。モデルの精度低下と引き換えに、もう少し計算コストが小さいモデルを用いてNPSを上げたほうが強くなる可能性が高いです。
model.predictを呼ぶたびにバッチサイズが異なる入力を与えても正常動作しますが、バッチサイズを変えた初回の実行は、2回目以降の実行より時間がかかるようです。例えば、Neural Engineでアプリ起動直後からバッチサイズ1024→1023→1024→1024という与え方をすると、処理時間が1.64秒→1.42秒→1.40秒→0.99秒というように変化しました。メモリの確保しなおしなどのオーバーヘッドが生じているものと思われます。そのため、サンプル数が少ない場合でも、ダミーデータで埋めて固定バッチサイズで実行したほうが良い可能性が高いとわかりました。さらに、CPUではバッチサイズを変化させた直後のほうが速いという不思議な挙動がありました。バッチサイズ8→8→1→1という与え方をすると、0.046秒→0.079秒→0.015秒→0.037秒という変化でした。なぜかバッチサイズ1にした直後が高速です。熱的な要因を疑って各実行の間を数十秒あけてみましたが変化はありませんでした。
列挙タイプのモデルも測定を試みましたが、Neural Engineで内部エラーが発生し実行できませんでした。CPUでは正常動作しましたが、範囲タイプの時と処理時間の差は（簡易測定の範囲では）見られませんでした。ちなみに、バッチサイズ1,2,4をEnumeratedShapesに与えて変換していても、その途中にあるバッチサイズ3の入力をするとエラーになります。自動的にパディングして大きいほうに合わせてくれるというような機能はありません。
以上のように、バッチサイズ可変モデルは一応作れますが、入力バッチサイズを変化させた直後の処理が遅いこと、不可解なエラーが発生することから使いづらいという結果になりました。今後、バッチサイズ16等で固定したモデルを生成してみて動作確認します。
※実験後にiPadOS 15.3にアップデートしました。範囲タイプのモデルを試した結果、速度、エラーとも本記事の内容と変化はありませんでした。
