TensorFlowのRust bindingにプルリクを送った話
この記事はRust Advent Calendar 2021 19日目の記事です。
ことの発端
この話は某ばんくしさんとのTwitterとのやり取りから始まりました。
元々の文脈は、僕がaxumとtch-rsでRustの画像認識APIを作るを見て、tch-rsをTFで置き換えたバージョンを作ってみたことでした。ただ、TFに置き換えて見ると、画像のエンコード/デコード周りをうまく扱うのが難しく、image crateに頼らざるを得なかったのですが、それに加えてRGBとRGBAの違いを解消するのにバッファを何度もコピーする必要がありました。上記のtwitterのやり取りがまさにそれです。また、似たようにところではOpenCVとの連携でCreating tensors from opencv dataというIssueも上がっていましたので、Rust bindingを使う人にとっては悩みの種であることが匂ってきます。
ちなみにTFだけでも標準的な画像のエンコード/デコードも(当然ながら)扱うことができます。例えば、画像を読み込んでTensorを構築するだけなら下記のように2, 3行で出来ます。
import tensorflow as tf
fname = "xxx.jpg"
buf = tf.io.read_file(fname)
img = tf.image.decode_jpeg(buf)
Rustでもこれを簡単に出来れば、上記の画像フォーマットやバッファのやり取りに関する問題を少しは軽減できるはずです。
Eagerモード対応がいる。
さて、上記のPythonのコードではTFでの操作はただの関数のように見えますが、これはTF2.x系でのお話です。「今は昔・・・」となりそうなのと僕自身も詳しくないので割愛しますが、デフォルトで関数のように扱えるようになったのはTF2.x系からEagerモードが標準になったおかげなのです。
Eagerモードをどうやって実装するよ
ところが、RustのTF bindingは基本的なC-APIに対するラッパーのみを提供しており、このようなEagerモードでの呼び出しは想定されていませんでした。
コードベースで言うと、Rustがラップしているのはこれ。
ところが、Eagerモードに関するAPIは別のヘッダーで定義されており、上記とは異なるAPIを使って構築しなければなりません。それがこっち(Zennのカードだと上記と区別付かないかもですが、eagerディレクトリのしたのc_api.hです)。
原理的な話をすれば、TFの操作は昔ながらのAPI(Define-and-Run style)でもEagerモード(Define-by-Run)でも同じことができますが、関数のように呼び出せたほうがシンプルで使いやすくてデバッグもしやすいはずです。
実は、こんな便利なモードを使えるC++ライブラリが一つありました（他にも有ったら教えて下さい）。
Eager APIを使用したライブラリ例：cppflow
cppflowはC++用のヘッダーライブラリなのですが、TFのC-APIをうまくラップして使いやすくしています。次のコードをPythonと見比べるとよく分かると思うのですが、とても似たような呼び出しになっています。
// Load an image
auto fname = std::string("image.jpg")
auto buf = cppflow::read_file(fname);
auto input = cppflow::decode_jpeg(buf);
これをRustでもやりたいわけです。お手本があるので簡単ですね(?)。これと、TFのリポジトリにはeager C-APIのUTコードもあります！これだけあればかろうじて何とかなりそうです。
(なお、これ以外に手掛かりになるようなドキュメントはありませんでした。ひぃ)
さっそく実験（そして超巨大PRの爆誕・・・
下調べを終えてやることが見えて来たので、いろいろ実験して何とか動くようになりました。ここで少し欲が出てしまい、本家にも取り込んでもらえば・・・という下心が芽生えました。
その結果、次のような超巨大PRが爆誕しました。
なんですか、+152k/-5kって。誰もレビューしてくれるわけないじゃないですか、アホなんですか。はい、反省してます。
それでも、なんとメンテナーの方がこれを時間掛けてチェックしてくれていたのです。ただただ頭が下がるばかり。
タスク整理
さて、こんな超巨大なままでは何にもできないのでいくつかに分割しました。
- eager C-APIへのバインディング
- 実際にはtensorflow-sysのサブcrateの修正
- Rustらしい呼び出しができるようにラップする。
- raw_opsの実装
大きくはこの3つをこなして行く必要があります。ただ1つめはbindgenを使うだけなのでほとんど作業はありませんでした(ref: https://github.com/tensorflow/rust/pull/328 )。2つ目はRustらしいところは曲者で、FFIとライフタイム周りでパズルをやりながら現在進めています。
ところでraw_opsとはなんぞや
3つ目について少し補足すると、先に例を上げたPythonとcppflowは、呼び出し方は似たような形になってはいるものの、その実体が異なっています。
上記の例で続けると、Python側のコードはioモジュールの中の関数io.decode_jpegを使っていますが、cppflowやRustでこれからやろうとしているものはraw_opsというモジュールの中にある関数[raw_ops.decode_jpeg]をそのまま利用できるようにしてあるだけです。言い換えると、raw_opsというモジュールでTFが実行できる全ての操作を列挙していて、ioモジュールなどはその操作をPython側で使いやすくラップしているのです(他のモジュールでもこの関係は当てはまります)。このため、cppflowでも今後のRustのほうでも、生の操作はできるもののPythonと同等にというのは現実的ではないことは頭の片隅に置いておくといいかもしれませんね。
これから
先ほどのステップ2.のRustらしい呼び出しができるようにラップする。というステップの半分くらいが終わったところでしょうか。最初の超巨大PRが発生してから3カ月程たっているので、2.と3.が終わるのも後3カ月くらいかかるかもしれません(クリスマスプレゼントには間に合いませんでした)。この先もうまく進むかは何ともですが、もし無事にマージされてリリースしてもらえたらぜひ使ってみて下さい。
Discussion