PyTorch Mobile

PyTorch 1.3 の一部。Caffe2 の名前が変わったのかと思ったら、ほんとに PyTorch の一部らしい。ということで久しぶりにサイトなどを冷やかす。以下わかったこと:

PyTorch は 1.0 あたりで TorchScript なるものが追加された。TorchScript は Python のサブセットを C++ で実装したバイトコードインタプリタみたいなもの。コード表現 (IR) は直列化もできる。Python と違って GIL もないし機能も少ないので速いしスケーラブルだと主張している。

TorchScript のコードを作る方法は２つある。一つ目は "tracing". モデルの実行をフックすることで、どんな計算を実行するのかを記録する。Autograd 用に計算を追跡する仕組みはすでにあるので、その横でついでに tracing  すればよい（割とひどいコードになっている)。ここには関数呼び出しみたいな概念がない。計算はぜんぶ「インライン化」される、つまり分岐もなにもない計算だけの IR が生成される。

IR をつくる２つ目の方法は Python AST からの変換。Python の AST を TorchScript の AST に変換し、それを IR に落とす。Python AST から変換すると関数呼び出しやループや分岐といった情報の保存された IR ができる。

Tracing と AST の住み分けはよくわからないが、Python 上のモデルがサポートされている AST のサブセットに収まらない場合は tracing を使うのだろうか？そんなんでいいのかよくあわからないけど。

Python サブセットのテキストをパースして TorchScript の AST を作る仕組みもあるという。エンドユーザには関係ないが、内部的に python の build-in を実装するのに使っているらしい。これかな？この規模ならパーサ書かなくても C++ の DSL でやればいいのでは、・・・。

PyTorch Mobile の実体はこの TorchScript である。AI 人材が Tracing なり AST 変換なりで TorchScript の IR をつくり、直列化して保存する。アプリ開発者は PyTorch Mobile の API でその IR をロードし、外からデータを渡して実行する。TorchScript (IR の実行部分) は C++ なので Python が要らない。ちなみに直列化フォーマットは Python pickle (のサブセット) で直列化した IR が zip (のサブセット) でパッケージされている。ぜんぶ C++ で書いてあるのにこれとは、さすが Python first を名乗るだけあるな・・・（なお最初は JSON をつかっていたらしく "LEGACY" とマークされた JSON ロード用のコードが残っている。

TorchScript と並行して CuPy にあたるレイヤも書き直しているっぽい。まえは ATen とかいうやつを使っていたが c10 というプロジェクトに置き換えようとしている、らしい。

あと  TorchScript はいちおう different-i-able ではある模様。ただどのくらいちゃんと動くのかは不明。inference だけなら backward path はいらないし。

以下感想とか。

TF との比較でいうと、TorchScript は PyTorch として AST/Graph を一級市民にするという意味で TF 的といえなくもない。ただ SavedModel と比べてふつうの言語っぽい IR を全面に押し出しているところが違うといえば違う。TF だと XLA とか MLIR みたいな IR のレイヤは SavedModel の下にある。

PyTorch の一族には Glow という中間表現およびコンパイラがいた。こいつがどうなったのかと覗いてみると、特殊な ML ハードウェア向けコンパイラインフラという TorchScript より一段下の位置づけになった模様。じっさい Glow のツリーには TorchScript IR をロードするフロントエンドみたいなコードがある。TorchScript 自体はネイティブコードへのコンパイルとかできないインタプリタなので、CPU でも E2E のネイティブコード生成をしたいなら Glow がやるのかもしれない。

PyTorch Mobile が TorchScript ベースだとすると少なくともモバイルの Caffe2 はさよならなのかな。モバイルにかぎらずプロダクションは Caffe2 を捨て TorchScript に移行する見通しなのだろう。そして Caffe2 がいらないならモデルの可搬性が必要なくなるから ONNX もいらないね。Glow も TorchScript IR を直接ロードしているし。

リサーチ志向の PyTorch が苦手なプロダクションの穴を埋めようと ONNX 連合に参加した他の ML フレームワークにしてみたらひどい話だけれど、これが大企業同士の戦争というものなのでしょう。

ふたたび TF との比較に戻ると、PyTorch は一つの TorchScript 実装でプロダクションをサーバからモバイルまでカバーしようとしているように見える。一方 TF はモバイル用に TF Lite を作った。TF Lite は自己完結したシンプルなインタプリタで、カーネルとかもぜんぶ TF とは独立して再実装している。TF はもともと C++ 中心で TFX Serving とかも Python なしに動かせるんだから、モバイルもそれでよかったのではという気がしないでもない。少なくともある時期までは Non-Lite な TF をモバイルで使っていたはずだから不可能ではないはずだけれど、どうして再実装したんだろうね。でかすぎたのかな？

TF Lite はモデルのフォーマットも flatbuffer で TF 本体と互換性がなく、toco というツールで SavedModel を変換しないといけない。Python 上の PyTorch から .pt ファイルを保存してそれを直接モバイルにロードできる PyTorch スタックと比べ moving parts が多く cumbersome に見える.

TF は単一開発元で E2E で全部もってるのが売りだったけれど、フルスタックな割に中身がバラバラ（特にモバイル）。PyTorch は Caffe2 だの ONNX だの Glow だのとバラバラだったのが、1.0 過ぎたあたりから統合が進み部品も揃ってきた。

そして朝に 1-2 時間くらいパラパラとコードをドキュメントを眺めただけでここに書いたような理解ができるくらいドキュメントも揃ってるしコードもコンパクト。実際どうなのかはさわってみないとわからないけど、触ってみたくなる空気はある。