少し前まで、MPSGraphBuilderというプロジェクトを地道に開発していました。
Core MLの.mlmodelファイルを読み込んで、MPSGraphを構築するというものです。
対応する層が限定的なものですが、おかげで、囲碁AIのLeela ZeroやKataGoのニューラルネットワークをMetal上で動かすことができました。
macOS Sonomaでmpsgraphtoolという、.mlmodelファイルを.mpsgraphpackageファイルに変換するツールが付属されるようになったので、MPSGraphBuilderの役目は終了しました。
この開発でドキュメントにない大きな知見を得たので、せっかくなので記録として残しておきます。
それは、Float16型のウェイトのMPSGraphを構築すると、AシリーズMシリーズではそれはApple Neural Engine上で動くということです。
MPSはMetal Performance Shadersの略で、MetalはGPUのドライバのブランド名なので当然GPU上で動くと思っていたのですが、MPSGraphのグラフコンパイラはデフォルトの最適化オプション(レベル1)では、使える時にはNeural Engineを使うコードを生成します。
すごくないですか？
Core MLモデルと違って、MPSGraphは動的に生成しやすいので、PyTorchやTensorFlowのバックエンドとしても利用されています。
つまり、PyTorch/TensorFlowで注意深くウェイトのデータ型をFloat16にすれば、少なくとも推論時にはGPUだけでなくNeural Engineも利用ということです。
この話、聞いたことがなくて もしかしたらどこかに落とし穴があってPyTorch/TensorFlowで使えていないのかもしれません。
実際、試しにPyTorchのresnet50をモデル、入力ともhalf()してもNeural Engineでは動きませんでした。
しかし、MPGGraphBuilderでMPGGraphがNeural Engineを使う条件が存在することは確認しました。
この話が広まって、Mac上でNeural Engineの活用がもっとカジュアルになれば幸いです。
