ChainerX のインストールと高速化の検証

インストール

まずここからは Python 3.6 で既に Chainer (Python) を GPU で実行できる環境を構築している前提で進める。また環境は基本的に $HOME/local の下にローカルインストールしており、そこにライブラリとかをガンガン入れていくことにする。

ChainerX を使いたい場合、Chainer インストール時に C++ のコンパイルが必要になる。GPU (cuda) を使いたいので、まず環境変数を設定する。

Installation — Chainer 6.0.0rc1 documentation

export CHAINER_BUILD_CHAINERX=1
export CHAINERX_BUILD_CUDA=1
export CUDNN_ROOT_DIR=/usr/local/cuda/lib64 # libcuda.so 的な奴の場所
export MAKEFLAGS=-j8  # 並列化する CPU の数

ここで、このサーバーで C++ を使ってなさすぎて環境が全然構築できていなかったので、cmake と g++ をローカルインストールした。ただし、ChainerX は gcc 6 以下しか使えないようなので、それに準拠して 6.5 を入れてみた。このへんは C++ のルーチンワークなので説明不要かもしれないが、馴染みのない方向けに。

CMake の最新は 3.14 のようなので、これを入れる。
https://cmake.org/download/

cd $HOME/local/src
wget https://github.com/Kitware/CMake/releases/download/v3.14.3/cmake-3.14.3.tar.gz
tar -zxvf cmake-3.14.3.tar.gz
cd cmake-3.14.3
mkdir build; cd build
../configure --prefix=$HOME/local
make -j8
make install

基本のインストールは以降全部一緒。gcc のローカルインストールは以下の記事が詳しかったので、これに従って m4-latest, gmp-6.1.2, mpfr-4.0.2, mpc-1.1.0 → gcc-6.5.0 とインストールした。

hatyuki.hatenablog.jp

gcc のインストールオプションはこんな感じ。

../configure --prefix=$HOME/local/gcc --with-gmp=$HOME/local/gcc --with-mpfr=$HOME/local/gcc --with-mpc=$HOME/local/gcc --enable-languages=c,c++ --disable-bootstrap --dsiable-multilib
make -j8
make install

これで準備完了したので、環境変数周りを調整してインストール。

export CC=$HOME/local/gcc/bin/gcc
export CXX=$HOME/local/gcc/bin/g++
export PATH=$HOME/local/gcc/bin:$PATH
export LD_LIBRARY_PATH=$HOME/local/gcc/lib:$LD_LIBRARY_PATH
export LD_LIBRARY_PATH=$HOME/local/gcc/lib64:$LD_LIBRARY_PATH
export LD_LIBRARY_PATH=$HOME/local/lib:$LD_LIBRARY_PATH
unset LIBRARY_PATH # なぜかエラーが出たので
pip install cupy --pre --upgrade --nocache-dir -vvvv
pip install chainer --pre --upgrade --nocache-dir -vvvv

これで、無事コンパイルができれば Chainer で ChainerX が利用できるようになる。

実験

ChainerX のインストールが完了したので、example に基づいて MNIST で実験をしてみる。

chainer/chainerx_cc/examples/mnist_py at master · chainer/chainer · GitHub

ChainerX のいいところは、従来の Variable を通さずにそのまま backward ができるところにある。しかし現在の Chainer のモジュールは非対応のものが多いので、独自関数を呼びつつ自分で custom loop を作成する必要がある。こういうコアな部分は天才エンジニアの方たちがいずれ解決してくれると思うので、 大きくいじらずに実験してみる。

従来の Chainer と ChainerX について、公平な比較は難しいが実験条件を決めておく。大きな違いとしては Chainer 側は chainer.iterators を使うことにする。あとは基本的に同じである。

今回は ChainerX のチュートリアルと examples にしたがって、以下のように Link, Model, Optimizer(SGD) を定義することにして、従来のソースコードから少しの修正で同じように呼び出せるようにしておく。ここもきっと天才エンジ(ry

ChainerX Tutorial — Chainer 6.0.0rc1 documentation

import numpy as np
import chainerx as chx

# chx Link
class Linear(object):
    def __init__(self, n_in, n_out):
        W = np.random.randn(n_in, n_out).astype(np.float32)
        W /= np.sqrt(n_in)
        self.W = chx.array(W)
        self.b = chx.zeros((n_out,), dtype=chx.float32)
    
    def __call__(self, x):
        x = x.reshape(x.shape[:2])
        return x.dot(self.W) + self.b
  
    @property
    def params(self):
        return self.W, self.b

class MLP(object):
    def __init__(self, n_units=1000, n_out=10):
        self.l1 = Linear(784, n_units)
        self.l2 = Linear(n_units, n_units)
        self.l3 = Linear(n_units, n_out)
  
    def forward(self, x):
        h = chx.relu(self.l1(x))
        h = chx.relu(self.l2(h))
        return self.l3(h)
  
    @property
    def layers(self):
        return (self.l1, self.l2, self.l3)

# chx Optimizer
class SGD(object):
    def __init__(self, lr=0.01):
        super(SGD, self).__init__()
        self.lr = lr
  
    def setup(self, model):
        self.layers = model.layers
        # require grad
        for layer in self.layers:
            for param in layer.params:
                param.require_grad()
        
    def update(self):
        for layer in self.layers:
            for param in layer.params:
                p = param.as_grad_stopped()
                p -= self.lr * param.grad.as_grad_stopped()
                param.cleargrad()

ミニバッチ1回の処理について、以下の5つの実行時間を比較してみる。

1. ミニバッチの取得
2. forward 計算
3. loss 算出
4. backward 計算
5. optimizer によるモデルの更新

実験はMNISTを使い、ミニバッチ100枚、epoch 数10回、時間は GPU で cp.cuda.Stream.null.synchronize() を実行して測定する。ついでに、loss と精度、全体の実行時間を比較する。実験結果は以下。

実際の重要な計算は同じ GPU で動いているので劇的に改善というわけではないが、それを差し引けば ChainerX が非常に高速で実行できていることが分かる。まだ無骨な印象が強いので、Trainer 周りも含めた今後の開発に期待している。

(おまけ) ChainerX C++ でそのままコンパイルするサンプル

GitHub に ChainerX を C++ そのまま書くサンプルがあるので、それも実際に動かしてみる。

chainer/chainerx_cc/examples/mnist at master · chainer/chainer · GitHub

このままでは僕の環境ではコンパイルできないので、ちょっと CMakeLists.txt をいじってみることにする。

まず、ChainerX が呼んでる外部ソース (.h のみで使えるやつ) gsl-lite, optional-lite をローカルインストールしてしまった(めんどくさいので)。

GitHub – martinmoene/gsl-lite: gsl lite – A single-file header-only version of ISO C++ Guidelines Support Library (GSL) for C++98, C++11 and later

GitHub – martinmoene/optional-lite: optional lite – A C++17-like optional, a nullable object for C++98, C++11 and later in a single-file header-only library

cd $HOME/local/src
git clone https://github.com/martinmoene/gsl-lite
cd gsl-lite
mkdir build; cd build
cmake .. -DCMAKE_INSTALL_PREFIX=$HOME/local
make -j8
make install

こんな感じで、$HOME/local/include にコピーされるので、そこを参照する cmake ファイルを書けばいい。あとは既にインストールしてある Python の中の ChainerX のライブラリを参照するように書けばコンパイル可能になる。(むしろこのサンプルはどうやって実行することを想定しているんだろう…)

set(ENV{CC} "$ENV{HOME}/local/gcc/bin/gcc")
set(ENV{CXX} "$ENV{HOME}/local/gcc/bin/g++")
# project
cmake_minimum_required(VERSION 3.14)
project( train_mnist )
set(CHX_DIR $ENV{HOME}/local/lib/python3.6/site-package/chainerx)
set(CHX_INC_DIR $ENV{HOME}/local/include)
add_executable(${PROJECT_NAME}
  mnist.cc
  train_mnist.cc
)
# include
set(INCLUDE_DIR
  ${CHX_INC_DIR}
)
target_include_directries(${PROJECT_NAME}
  PUBLIC "${INCLUDE_DIR}"
)
# lib
set(LIBRARIES
  ${CHX_DIR}/libchainerx.so
)
target_link_libraries(${PROJECT_NAME}
  PRIVATE ${LIBRARIES}
)
target_compile_options(${PROJECT_NAME} PUBLIC -O2 -Wall)
target_compile_features(${PROJECT_NAME} PUBLIC cxx_std_17)

こんな感じでいいかな。あとは以下のように実行できる。

../mnist_py/download.sh # 同じ examples にダウンロードスクリプトがある 
mkdir build; cd build
cmake ..
./train_mnist --data ../mnist --device cuda:0 # 引数に ChainerX (cuda:0) を指定

実行時間は 21.3 秒だったので、C++ ベタで使ったほうが Python でマネージするよりちょっぴり早くて(当たり前とはいえ、Python そのままでも結構早いということが分かる) ついでに C++ でソースコードの隠蔽もできる。

今回遊んだコードは以下に置きました。

github.com

WE ARE X（字幕版）

• 発売日: 2017/12/13
• メディア: Prime Video
• この商品を含むブログを見る