ChainerX移行ガイド

chainerx 対応：コード変換部分早見表

とうとうChainer version 6.0.0 がリリースされましたね。本記事ではこれまでの Chainer v5 以前までに書かれていたモデル訓練用コードを、 Chainer v6 で導入されたChainerXへ対応させるための変更点についてまとめます。
（上記コードの変換をしなくても今まで通り v5 のコードを v6で動作させることは可能です。）
本題に飛びたい方は v5 のコードを v6 のChainerX対応コードに変換する からご覧ください。

ChainerX とは

先日Chainer version 6.0.0 が正式リリースされました。

• Released Chainer/CuPy v6.0.0

#Chainer/#CuPy v6.0.0 をリリースしました。変更点の概要はブログをご覧ください。 https://t.co/1k7VXirOFC　https://t.co/Rf7GxJCZyKhttps://t.co/leNCHR6bdU

— ChainerJP (@ChainerJP) May 16, 2019

V6で追加された機能で注目したいのは何といってもChainerXです。
ChainerXはC++で実装されており、
- 多次元配列の計算 Operators：numpy arrayのように使える
- Device management　　　　：さまざまなDevice (CPU/GPUなど) を切り替えて使える
- 自動微分機能 　　　　　　　：chainer Variableが行っていた、微分機能

の機能を合わせもつ Array を提供しています。

これまで numpy or cupy という CPUかGPUか、→他のDeviceは考えていないというコードを書いていましたが、ChainerXの導入により将来新しく出てくるDeviceに対してもアプリケーション側のコードは同じまま、計算を扱えるようになりました。

つまり、CPU だろうがGPUだろうが、理論的には "Backendさえ追加で実装されれば" (計算で使用する関数など)、 IOTに使用される格安チップで使われるFPGA、スマホなどで使われる Arm 専用などなどにも対応できる仕組みということです。

ChainerXに変えることによる大きなメリットは以下二つです
・高速：C++で書かれているため、python のOverheadを少なくすることができます。※
・hetero device への対応：Backend さえ実装すれば、組み込みなどへの対応も可能です！

※ ただし、現状はBLASなどの実装が最適化されていない関数もあり、絶対に numpy/cupy より速い！とは限らないようです。

より詳しい説明は以下のslideshareやブログをご覧ください
- ChainerX and How to Take Part
- N次元配列の自動微分をC++で実装したChainerXをリリース。Chainer v6（β版）に統合し、計算パフォーマンスを向上

インストール方法

• https://docs.chainer.org/en/latest/chainerx/install/index.html

ChainerXはDefaultではBuildされません。
Linux では以下のように環境変数を設定してからChainerをインストールすることで、ChainerXもBuild されます。

$ export CHAINER_BUILD_CHAINERX=1
$ export CHAINERX_BUILD_CUDA=1    # Only when you have GPU and CUDA installed
$ export MAKEFLAGS=-j8            # Using 8 parallel jobs.
$ pip install chainer

Dockerfileを使いたい場合は、以下のように書くとよいでしょう。

FROM nvidia/cuda:10.0-cudnn7-devel-ubuntu16.04
...

# Install from source is expected, to align cudnn version between cupy & chainerx
RUN pip install cupy==6.0.0

# chainer & chainerx
ENV CHAINER_BUILD_CHAINERX 1
ENV CHAINERX_BUILD_CUDA 1
ENV CUDNN_ROOT_DIR=/usr/include
ENV MAKEFLAGS -j8
RUN pip install chainer==6.0.0
#RUN git clone https://github.com/chainer/chainer && pip install -e chainer

# To use ChainerMN
# RUN pip install mpi4py

GPUを使わない場合は cupy のインストールと、 CHAINERX_BUILD_CUDA, CUDNN_ROOT_DIR の設定の行を抜かしてください。

v5 のコードを v6 のChainerX対応コードに変換する

いよいよ本題です。既存のChainerコードをChainerXに対応させるためにコードのどの部分を変換すればよいかを紹介します。

一番BasicなOfficial Example である train_mnist.py を、 v5のコードとv6のコード で比較してみればその違いはすぐに分かります。

一番の大きな違いは以下の部分です。

~v5: CPUかGPUかを、args.gpu として int で指定。 GPUの場合のみ Deviceへ送る処理を書く。

parser.add_argument('--gpu', '-g', type=int, default=-1)
...

# args.gpu は int. -1 がCPU、0以上が使用するGPU IDを表す。
model = L.Classifier(MLP(args.unit, 10))
if args.gpu >= 0:
    # Make a specified GPU current
    chainer.backends.cuda.get_device_from_id(args.gpu).use()
    model.to_gpu()  # Copy the model to the GPU

v6: どの deviceをつかうかを、args.device として str で指定。 任意の device へ送る処理を書く。

parser.add_argument('--device', '-d', type=str, default='-1')
...

# args.device は str となり、Deviceを文字列で表すように。
device = chainer.get_device(args.device)
model = L.Classifier(MLP(args.unit, 10))
model.to_device(device)
device.use()

基本的にはこの部分を変更するだけでOKです。
上記コードの変換をするとChainerXに対応した動作ができるようになるということなので、変換をしなくても今まで通り v5 のコードを v6で動作させることは可能です。

device.use() は必須ではなく省略可能のようです。
また有効範囲としては、thread local な機能なので、別スレッドでは効果がないということに注意してください。

以下、コード上は変化ありませんが、関数の引数が変化している部分です。
これまで device として -1 or 0 といった int を引数としていたところが、 chainerのDevice instance を受け取るように変化しています。

• Updater
• Evaluator

~v5

device = args.gpu  # type: int
updater = training.updaters.StandardUpdater(
        train_iter, optimizer, device=device)

trainer.extend(extensions.Evaluator(
    test_iter, model, device=device))

v6

device = chainer.get_device(args.device)  # type: chainer._backend.Device
updater = training.updaters.StandardUpdater(
        train_iter, optimizer, device=device)

trainer.extend(extensions.Evaluator(
    test_iter, model, device=device))

また、converter をカスタマイズして使っている方はその部分も変更が入っています。

実行時の option の指定方法

これまで、
--gpu -1　→ CPU, numpy
--gpu 0 　→ GPU, cupy

と int 型を指定してCPU or GPUを切り替えていたののが、

--device -1　 → CPU, numpy
--device 0　 → GPU, cupy
--device native → CPU (native backend), chainerx
--device cuda:0 → GPU (CUDA backend), chainerx

というように str 型を渡すことで任意のDeviceを使えるような設計に変わっています。
'-1' や '0' といった int に変換できる文字列を渡した場合には これまでどおりの挙動と同じく numpy/cupy が使われます。

今後新しいDeviceに対しても Backend が実装されれば、アプリケーション側のコードは同じまま、
--device custom_device:0 などと指定することで動かせることができるようになっています。

まとめ

以上をまとめると冒頭に張った図のように変換するだけでChainerXを動かせるということになります。

chainerx 対応：コード変換部分早見表

時間がとれれば次は少しコードリーディングや device class/backend 機構について追記したいと思っています。

皆さんもぜひChainerXを使ってみてください。
chainer レポジトリから、コードをpull して /examples/mnist で
python train_mnist.py --device cuda:0
などと実行してみれば動かせるはずです！