読み書きプログラミング

日常のプログラミングで気づいたことを綴っています

Leela ZeroのウェイトをCore MLで計算させる

AZ.jsをSwiftに移植しました。

その際に、Leela ZeroのウェイトをCore ML上に移す作業があったのですが、ハマり所があったので、備忘録です。

Leela Zeroのウェイトは、leela-zeroレポジトリのtraining/tfフォルダ以下のスクリプトを使って、TensorFlowに読み込むことができます。なので、後はtfcoremlで.mlmodelに変換すれば完了です。
github.com

以下、ポイントは2点ありました。

  • tfcoremlはテンソルのフォーマットにNHWCしかサポートしていません。レポジトリのtfprocess.pyは(GPU用に)NCHWフォーマットなので、NHWCフォーマットに対応させました。
  • Core MLは、現在、モデルの実装が限定されているようです。無次元+チャンネルの場合バッチあり、2次元+チャンネルの場合バッチなしの2つです。1次元+チャンネルの場合2次元+チャンネルの1次元を要素数1にして代用するようです。なので本当の意味での1次元+チャンネルはサポートされていません。Leela Zeroは1次元+チャンネルを内部でreshapeする構造なのでこのままではCore MLに合わせることはできませんでした。なのでLeela Zeroのモデルを、内部のreshapeを取り払って2次元+チャンネルの入力に変更しました。

ELF OpenGo v1のウェイト(20b, 224c)を変換してみたところ、私のiMac(27インチ ,Late 2012, i5)で推論が1回170--220ms程度でした。同じマシンのChrome(webgl)でWebDNN上のものが360ms前後なので倍ぐらいのスピードですか。Core ML(tfcoreml?)は量子化をしているようで出力の数字が荒っぽくなります。