転職してからやってること

転職してからやってるプロジェクトについて何か書いてみようかと思います。

https://github.com/pfnet-research/chainer-compiler/

公式ぽいブログも書きました。

https://research.preferred.jp/2019/01/chainer%E3%83%A2%E3%83%87%E3%83%AB%E3%81%AE%E3%81%95%E3%82%89%E3%81%AA%E3%82%8B%E9%AB%98%E9%80%9F%E5%8C%96%E3%80%81%E3%83%87%E3%83%97%E3%83%AD%E3%82%A4%E3%81%AE%E7%B0%A1%E4%BE%BF%E5%8C%96%E3%80%81/

あまり綺麗にまとまった文章を書ける気がしないので、つらつらと時系列で思い出して行こうかと思います。

転職前、 define-by-run とかについて考えていたこと

まだ前職にいた時、転職活動中に、 ChainerX の存在や、 Python からコンパイルして実行できると良いなと思っている、みたいな話を聞いたような気がします。これは結構、具体的かつ無茶ぶり感がある話で、まあ楽しそうだなと思ったし転職先を決める理由になったような気もします。

機械学習、どんな計算してるかというとおおむね行列のかけ算ばかりしている、とかよく説明します。とはいえ非線型な要素も必須だし、計算結果をどうつなぐか、みたいなのでなんかうまくいく方法が発見されたり、ブロックの接続をこねくり回してるような感覚があると思います。

で、そういう接続をごにょごにょする機械学習フレームワーク、色々あるわけですが、 Chainer 以前は、ブロックの接続を作る DSL があって、接続ができたら、誤差逆伝播の計算を考えて、ハイ実行、と実行する形式だったと。 tf.add(A, B) とかすると、 A + B を実行するんじゃなくて、足し算するための計算グラフ内のノードを返す、みたいな感じ

で Chainer は、 chainer.functions.add(A, B) とかすると、そのタイミングで実際に実行する。実行すると同時に、誤差逆伝播する時に必要なグラフも作っておく。コンパイラに対するインタプリタという趣きで、問題が起きた時に問題のある行で止まるので、デバッグがしやすい、分岐などを含むようなフレキシブルなモデルを普通に記述できる、などのメリットがあります。ただ、欠点が2つあります

  • Python 無しのデプロイができない。計算グラフを作る DSL という形式であれば、そのグラフをエクスポートしてしまえば Python 無しのアプリから使えるわけです。
  • 速度的にオーバヘッドがある。これは、 CUDA の関数実行が遅延されているおかげで、単純なモデルでは GPU 律速になって Python+CPU でグラフ作ったりしてるコストは隠れることが多いです。 cupy てセンスいいものだなあとつくづく思います。が、自然言語モデルや、小さいモデルだと CPU 律速になったりするし、 GPU が高速化し続けてることも相対的に問題を大きくしていっています。

ですが、まあ、研究者のイテレーション速度の方が重要だと考えていて、 define-by-run 良いなあと思っていました。ループがあるモデルを扱ってたことによるバイアスもあると思いますが。

そんな時になるほど!と思うものが現われました。 Swift for TensorFlow というやつです。 Swift for TensorFlow は、 define-by-run 、つまりモデル定義を実行時に行なうことが重要なのではなく、モデルがデータでなく、コードで定義されていることが重要なのだ、と主張しました。当時、 tf.while_loop じゃなくて Python の普通のループを使わせてくれ〜と思っていたので、このコンセプトが大変気にいりました。この図もなるほどーという感じでしたね

https://raw.githubusercontent.com/tensorflow/swift/master/docs/images/AutomaticDifferentiation-Approaches.png

グーグルの作ったもので Native Client の次くらいに好きなものといえると思います。つまり成功する気がしないのですが。 Swift 本当に使うようになるのか??という普通の疑問がまああるわけです。

で、そういうタイミングで、転職活動中に、具体的にどういう計画かよくわからないけど、 Python コンパイルしよう、という話だったので、本当にできるのかはよくわからないけど、それは関われるものなら挑戦してみたい話だな、と感じました。

このへんの DL フレームワーク話は 退職直後、入社前の間の時期にも書いたりしました

入社してみた

入ってみると、何も決まってないから、なんか IR をデザインして、 Python をそれに変換して、 ChainerX で実行するのやってくれ、という趣旨のことを言われたと思います。もうちょっと具体的に計画や実装が既にあるかと思ってたので、ずいぶん無茶振り来たなぁ、と思ったように思います。同日にインターンの人が入ってきて、その人には Python を ONNX という業界標準のモデルシリアライゼーションフォーマットに変換するタスクをやってもらおうと思っている、という話も聞きました。

いや似たような Python から変換するタスクを別々にやってどうすんねん、ということで、 ONNX をざっと眺めたところ、これを拡張したら十分にやっていけるだろう、と感じたので、インターンの人には Python => ONNX の部分、私の方は ONNX を実行するものを ChainerX を使って作る部分、と分担しようと考えました。これはなかなか良い決断だったんじゃないかなと思っています。実のところデザインとか全く自信が無いので、 IR とかうまく作れる気がしなかったし、ある程度考えられたデザインにタタ乗りできたのはおいしかった。あと、 ONNX のエコシステムまわりの他の人たちが作ったツールが使えたり、社内の ONNX 関係の別のツールとの協調とかも視野に入るのも良かったです。

このあたりで ChainerX についての理解も深めていきました。 Chainer は行列計算とかの重い計算は numpy/cupy というのに投げてるわけですが、 numpy とコンパチな ChainerX array というのを C++ で作って、 C++ 側に、従来 Python で実装されていた誤差逆伝播を持ってくる、というコンセプト。

前章でデプロイができない、速度が遅い、という問題に触れましたが、 ChainerX はこれらの問題を大幅に軽減するものです。 ChainerX array は C++ で書かれているので、 Python 無しのデプロイの道が開かれてくるし、速度が遅い一番の原因だったらしい、誤差逆伝播Python で書かれているという問題もなくなるのでした。

作ってみた

とりあえず世の中にある、ある程度大きいネットワークを動かせるようにしていきました。 ChainerX に基本的に必要な関数があるので、基本的には ONNX と ChainerX 関数を対応づけていくだけです。最初は ChainerX を使う C++ コードを生成していたのですが、テストのイテレーションが遅くなるのを嫌って、 ChainerX の関数を次々呼んでいくだけ、みたいな VM を定義して、それを使うようにしました。 ResNet50 くらいは割とすぐ動いたような気がします。

時系列は覚えてないですが、このへんで誤差逆伝播のグラフを生成するコードもでっちあげたのではないかと思います。最適化やらなんやらを実装したいとすると、計算グラフを編集するライブラリとして機能できる必要がどうしてもあるので、そのあたりの proof-of-concept として、自動微分ができる、というのは最低限の機能担保という感があったのでした。

そうこうしているうちに、インターンの方が同じく ResNet50 くらいのシンプルなネットワークを ONNX 化できるようにしてくれて、実際にそのグラフを入力としてトレーニング用のグラフを生成して ImageNet という、 224x224 の画像データを 1000 種類に分類、みたいな使ってトレーニング、みたいなことができはじめていたように思います。

このへんからは時系列曖昧ですね

EspNet を end-to-end で通そうとする日々

EspNet というのは、音声認識音声合成、その他もできるツールキットで、これのまあまあ複雑なコードパスを通そうと努力してました。とりあえず、ツールチェインの上から下までで、ループなどをちゃんと扱えますよ、という状態にしたいという気持ちがあったので、これをターゲットとした感じでした。 Conv + NStepLSTM + CTC + デコーダーは attention ついてるからループ回して LSTM やる、みたいな感じになってて、色々なものに対応しなければならないのです。

元々の ONNX の Loop や Scan といったオペレータは、固定長の ndarray 的なものの入力が入ってくる前提だったので、これは可変長の Python list とかが扱えるようにしたいなあ、と考えて、ここは結構大幅に拡張した部分でした。

インターンの人が、無茶ぶりにも関わらず、まあまあややこしいループなども動くように作ってくれていた感じでした。実際、ループ動かすための仕様みたいなのもキチンと決めず、かなりダメダメなメンターだったはずなのに、ちゃんとそれなりにそれぽかったのでとても感謝でした。それを引き継いで EspNet で必要な、かなりややこしい変換などをできるように努力していきました。 __init__ で self.x = None しておいて、後で if self.x is None: self.x = ... とかする、遅延初期化みたいなのがなかなか面倒でした。

それが通ると、次はループや if の微分というか誤差逆伝播を頑張って、あとは Python list まわりも逆伝播実装して、これらもなかなか大変でした。

このあたり、ループをちゃんと動かすとか、 Python の list を扱える、とかにこだわっていたのは、 Chainer のフレキシビリティを損ないたくないなあ、と思っていたことがあります。 dynamic shape とか可変長 list 扱うとかすると、どうしても最速が出ない局面は出てくるとは思うのですが、とりあえずフレキシビリティの方が重要かなと。あと dynamic shape を扱いつつも、 static に決まる部分は決める、みたいな方針でやってるので、次元が決まったところは高速化…みたいなこともできるのではないかと思っています。実装が追いついてないですが、 TVM のコード生成使う実験とかは shape がわかってる前提でやっていました(というか TVM は dynamic shape に関して結局なんか計画あるのだろうか)。

EspNet をまともな速度にしようとする日々

動くようになった最初はなんか、素の Chainer より圧倒的に遅い!みたいな状態でした。で…何したんでしたっけ。ちゃんと cuDNN の LSTM を使うようにしたり、 elementwise op の fusion したり、定数伝播やったり、などなど。 Python からグラフ生成してると、なんかゴミみたいなノードができまくるので、定数伝播みたいなのが重要なのですよね…

でまあ、まともな速度になって、 CPU heavy なケースでは Chainer に勝つなあという感じになりました。ただ CTC とかいう処理を実装してなかったので、実際に end-to-end で訓練することができなくて、しかしそんなの実装するのかったるいので、 Python インターフェイスを追加して、モデルの一部分だけを chainer-compiler で処理する、みたいなこともできるようにしました。

でまあそれぽくなりました。 Python コードに多少の、気持ち悪い変更が必要なのが、なんともはやな部分ではあるのですが…

TVM で遊ぶ

年末あたり、 TVM で遊びました。 LSTM は、会社としては直接そんなに重要でないということもあり、ちゃんと Conv とかを良くできるポテンシャルはそなえておかないとかなあ…と。

結果としては、色々面白いけど、 PythonC++ 行ききしまくりのコードがつらいというのがまず感想でした。あと、特に普通の Conv とかだと、伸びしろがそもそもあまりなくて、投下労力に対して得られるものがなあ…という気持ちもありました。

で、とりあえず保留することにしました。いずれにせよ、フレキシビリティという観点などから、自力でカーネルを生成できるオプションは中期的にはあるべきだと考えていて、どこかの段階で戻ってくることになると思っています。

オープンソース

これやること無限にあるから、一人でやってると崩壊以外見えねえなあ…と思ってたので、人が欲しいです、みたいなことをわめき始めてました。これグーグルだと、人が増えるかあるいはお前のそれやってるそれ意味ないからやめろや、て言われる二択でわかりやすかったのですが、小さい会社でみんなそれぞれ忙しくて、人も増えないがお取り潰しにもならない、という雰囲気でした。

で、どうも社内で取ろうとするんじゃなくて、外から拾ってこい、という文化だと学びました。確かに、個々人で拾って来させるってのは、なるほどベンチャーの成長方法として理に適っているなあ、と思いました。というわけでツイッターでわめいたりしてみました。

ただ、なんかやっぱコミュニケーションがしづらいということを思っていて、ユーザ向けリリース、という感じではなくて、こういうことやってますよ、仲間募集、ということでオープンソース化したいと思って、しました。とはいえ、それなりには、色々整える必要があって、それなりにはめんどうでした。

これから

適当に風呂敷を広げて、結構広い範囲で実装してみたけど、とりあえず、という感じの仮立て付けぽい状態の部分が多い状態で、まだまだやることが無限にあるなあ、と思っています。

基本的に、 chainer-compiler にこだわらず、 Chainer まわりで静的なグラフを扱うシーンのエコシステムを、全体として整えていきたいな、と思っています。最近ちょっと onnx-chainer に手を出してみたりもしていたり、とか。このへんキチンとしてないことがデプロイの大変さみたいなことにも直結してしまっているので。

不完全な部分がそこらじゅうにあって、どこから手をつけても良いような感じだし、やっぱ社内需要を満たすところを優先するのが良いよねえ、と思っているので、社内需要満たす努力をしつつ、汎用的なツールチェインとして少しずつ育てていけると良いなぁ、とか思っています。

あと例のごとく一緒に働きたい人募集中です。うまくいくかは不明ですが、ある種の人には、楽しくはあるんじゃないかと思います。少なくとも僕は大変楽しんでいます。

なにかあれば下記メールアドレスへ。
shinichiro.hamaji _at_ gmail.com
shinichiro.h