本日は（も）

カレンダー空いているのでChainerネタ書いてみます。
２，３日目は ChainerX という未来を git clone Chainerのリポジトリ を導入ることでできました。
１２月３日に v6.0.0b1 のリリースアナウンスがあったのでリポジトリを落とさなくても

$ pip install --pre chainer

によって未来へタイプスリップできるようになりました。

同時に 5.1.0 のバージョンもできたのでV５で導入された　Static Subgraph Optimizations　を使ってみようと思います。

概要

オープンソースの深層学習フレームワークChainer および汎用配列計算ライブラリCuPy の最新版となるv5をリリース

によれば

学習を通して変化しない静的グラフをキャッシュすることで、その部分の計算やメモリ使用を最適化し、学習速度を20～60%高速化

とのことです。

v5.0.0rc1 のリリースノートでは

Static subgraph optimization feature has been introduced. The CPU (Python) overhead of graph construction and traversal in backward is removed with it.

となっています。計算グラフの構築におけるPythonのオーバヘッドっでそんなに重いんですね・・・。

Advanced graph optimizations/transformations are not implemented yet, so currently it only reduces the CPU overhead. We will consider adding more sophisticated graph-level optimizations to improve the GPU utilization as well as further reduce CPU overhead.

要するにもっと改善の余地はあるようなので、今後のアップデートによって学習の高速化が実現できるかもしれません（わきゅわきゅ）。

Chainer における
深層学習の高速化
Optimizing Deep Learning with Chainer

の資料を見るといろいろなノウハウで学習速度を速めている様子がわかります。

Install

未来へタイムスリップしたお客様は念の為アンインストールしておきます

$ pip uninstall chainer cupy

そして再度インストールします

$ pip install chainer==5.1.0 
$ pip install cupy-cuda90==5.1.0

単に pip install cupy でもいいですがビルドが走るので cupy-cuda90 ないしは cupy-cuda92 などでもよいでしょう。くわしくはこちら

https://docs-cupy.chainer.org/en/v5.1.0/install.html#install-cupy

$ cat /usr/local/cuda/version.txt

で各自のバージョンを確認しておくとよいでしょう。

使い方

もう悩む必要はありません！ドキュメントを眺めればほら簡単（テレビショッピング風）

Static Subgraph Optimizations: Usage

To enable static graph optimizations, it is only necessary to add the @static_graph decorator to a chain’s call() method. We will now show how the Chainer MNIST example can be modified to use this feature. The modified version with static subgraph optimizations is located at chainer.examples.static_graph_optimizations.mnist.

つまり

import chainer
from chainer import static_graph

class BOKUNO_KANGAETA_SAITUYO_NETWORK(chainer.Chain):
    def __init___(self):
        super(BOKUNO_KANGAETA_SAITUYO_NETWORK,self).__init__():
        with self.init_scope():
            #write some links
    @static_graph
    def __call__(self,x):
        #write something

のようにすれば良いみたいですね。

デコレートについては　Pythonのデコレータを理解するための12Step などが参考になります。

試すコード

ChainerのリポジトリのExample　にもありますが、ここではChainerを勉強していた時に書いていた古き良き時代のCNNを用いたMNSIT学習コードを改造したものを使ってみます。

cnn_mnist.py

import argparse

import chainer
import chainer.functions as F
import chainer.links as L
from chainer import training, datasets, iterators, optimizers
from chainer.training import extensions
import numpy as np
from tqdm import tqdm

major, _, _ = chainer.__version__.split(".")
MAJOR = int(major)
if MAJOR >= 5:
    from chainer import static_graph
else:
    def static_graph(func):
        """
        dummy decorator to keep compatibility between Chainer v5 and v4
        """

        def wrap(self, *args, **kwargs):
            return func(self, *args, **kwargs)
        return wrap

DEVICE = 0
BATCH_SIZE = 32


class MNISTCONV(chainer.Chain):

    def __init__(self, use_static):
        super(MNISTCONV, self).__init__()
        ksize = 3
        hidden_ch = 8
        self.use_static = use_static
        with self.init_scope():
            self.c1 = L.Convolution2D(1, hidden_ch, ksize=ksize, stride=1, pad=ksize // 2)
            self.c2 = L.Convolution2D(hidden_ch, hidden_ch, ksize=ksize, stride=1, pad=ksize // 2)
            self.c3 = L.Convolution2D(hidden_ch, hidden_ch, ksize=ksize, stride=1, pad=ksize // 2)
            linear_size = 28 // 2
            self.l1 = L.Linear(linear_size * linear_size * hidden_ch, 10)

    def _forward(self, x):
        h = self.c1(x)
        h = self.c2(h)
        h = self.c3(h)
        h = F.relu(h)
        h = F.max_pooling_2d(h, 2)
        h = self.l1(h)
        return h

    @chainer.static_graph
    def static_forward(self, x):
        return self._forward(x)

    def forward(self, x):
        if self.use_static and chainer.backends.cuda.get_array_module(x) != np:
            return self.static_forward(x)
        else:
            return self._forward(x)

    def __call__(self, x, t=None, train=True):
        h = self.forward(x)
        if train:
            loss = F.softmax_cross_entropy(h, t)
            chainer.reporter.report({
                'loss': loss,
            }, self)
            return loss
        else:
            return F.softmax(h)


def train(model):
    if DEVICE >= 0:
        chainer.cuda.get_device_from_id(DEVICE).use()
        chainer.cuda.check_cuda_available()
        model.to_gpu()
    train, test = chainer.datasets.get_mnist(ndim=3)
    train_iter = iterators.SerialIterator(train, BATCH_SIZE, shuffle=True)
    test_iter = iterators.SerialIterator(test, BATCH_SIZE, repeat=False, shuffle=False)

    optimizer = optimizers.SGD()
    optimizer.setup(model)

    updater = training.StandardUpdater(train_iter, optimizer, device=DEVICE)
    trainer = training.Trainer(updater, (5, 'epoch'), out='result')
    trainer.extend(extensions.Evaluator(test_iter, model, device=DEVICE))
    trainer.extend(extensions.ProgressBar())
    trainer.extend(extensions.LogReport())
    trainer.extend(extensions.PrintReport([
        'epoch', 'elapsed_time',
        'main/loss', 'validation/main/loss',
    ]))
    trainer.run()

    chainer.serializers.save_npz('result/mnistconv.npz', model)


def predict(model, use_ideep):
    chainer.serializers.load_npz('result/mnistconv.npz', model)
    train, test = chainer.datasets.get_mnist(ndim=3)
    counter = 0
    acc = 0
    if use_ideep and chainer.backends.intel64.is_ideep_available():
        model.to_intel64()
    with chainer.using_config('use_ideep', 'auto'):
        for t in tqdm(test):
            counter += 1
            x, ans = t
            result = model(np.array([x]), train=False).data[0]
            if ans == np.argmax(result):
                acc += 1
        print(acc / counter)


def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--static', action='store_true', help='decorate static_graph on training')
    parser.add_argument('--ideep', action='store_true', help='use ideep on prediction')
    args = parser.parse_args()
    return args


def main():
    args = parse_arguments()
    klass = MNISTCONV
    use_static = args.static
    use_ideep = args.ideep

    model = klass(use_static)
    train(model)
    # initialize
    model = klass(use_static)
    predict(model, use_ideep)


if __name__ == '__main__':
    main()

このコードの使い方

$ python cnn_mnist.py

でいけます。ただし、この場合は Static Subgraph Optimizations の機能が使われないようになっています。

実行例を下記に示します。

epoch       elapsed_time  main/loss   validation/main/loss
1           13.5793       0.332974    0.149877              
2           25.9988       0.126915    0.0866168             
3           38.7518       0.0898346   0.0763928             
4           51.3052       0.0743761   0.0750592             
5           63.7166       0.0664786   0.0607661             
100%|██████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [00:27<00:00, 361.07it/s]
0.98

Static Subgraph Optimizations　を有効にするには --static をつけます:

$ python cnn_mnist.py --static

end of build_schedule()
Creating new backward schedule...
end of build_schedule()
epoch       elapsed_time  main/loss   validation/main/loss
end of build_schedule().......................................] 19.20%
end of build_schedule()#####################################..] 96.00%
1           12.0856       0.334284    0.135471              
2           23.1138       0.118735    0.0932434             
3           34.2854       0.0838322   0.0617251             
4           45.5717       0.0688249   0.054637              
5           56.6426       0.0596717   0.0525059             
100%|██████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [00:25<00:00, 388.59it/s]
0.9823

６３秒かかった学習が５６秒にへりました。約1.12倍はやくなりました。やったね！

うまくいかないこともある

まだ experimental な機能なのでうまく行かないケースもあります。今回は極めて単純なケースなのでうまく行きましたが、現段階ではいたずらに @static_graph をつけると全てうまく行くというわけではなさそうです。たとえば、今回のサンプルコードで

__call__ にデコレートすると失敗します。

リリースノートにあるようにこれらの機能を使うには forward のコードの書き方をそれを安全に運用できるように書き方に注意しないといけないようです。

By applying @static_graph decorator to functions or methods (typically it is the forward method of a chain), you can let Chainer cache the computational graph collected at the first call and reuse it from the subsequent calls. To use this feature safely, your define-by-run code must always perform the same computations each iteration.

できない例たち

ImageNet の Example にあるようなGoogLeNetの __call__　関数をデコレートしても end of build_schedule() メッセージが大量に出てきてメモリーオーバします。
MobileNetV2 で遊んでいると一応デコレートする前はちゃんとロスが収束していたのに、つけるとValidationが学習回しても減らないみたいなのはあります.(どのレイヤーが悪さをしているのかまだ切り分けできていない）
5.0.0 のときは chainer.links.BatchNormalization で use_gamma=False を指定したレイヤーを推論で通すと実行時エラーが起きる(再現コードなくしたのでなんとも言えない)

推論ではつかえるか？

原理的にはYesなはず。自分が書いたコードでGPUを用いた推論のスピードが向上するというケースもありました個人的にはこれだけでも嬉しいです。

ただし、今回の例では（コードを変更して推論時にも有効となるようにすると）うごくっちゃうごくんですが毎回　end of build_schedule() が走るのでうまくできていない様子がわかります。このプリントオプションなんとかしたい。

ideep と併用できる？

少し期待しましたが、だいたい失敗します。上のコードをいじると失敗させることができます。

デコレートすると Chainer V4 でコードが動かない(そりゃそうだ）。

V４以下では機能がないのでそりゃそうなんですが、V５が出たばかりで、V4で学習、推論ができた資産を云々という理由で実行環境として V4 と V5 で両立したいというニーズはあると思います。V5でこの機能をパーフェクトにつかえるシチュエーションがあるのにV4との互換性を意識して実装しない（あきらめる）みたいなのはちょっと悔しいですよね。

回避策（たとえば）

上のサンプルコードでも書いていますが、

major, _, _ = chainer.__version__.split(".")
MAJOR = int(major)
if MAJOR >= 5:
    from chainer import static_graph
else:
    def static_graph(func):
        def wrap(self, *args, **kwargs):
            return func(self, *args, **kwargs)
        return wrap

のようにしてChainerのメジャーバージョンに応じて static_graph の意味を変えるという方法があります。V4ではダミーのデコレータを定義する方法がありますね。

fowardの方も使用条件に応じて static_graph でデコレートされたところを通すか通さないかの条件分岐をつくると回避できます。たとえば、学習時はStatic Subgraph Optimizationsを活用したい。でも（何かしらの理由で）推論時は ideepで高速化させたいのときは

    def _forward(self, x):
        # ここが実質の本体

    @chainer.static_graph
    def static_forward(self, x):
        return self._forward(x)

    def forward(self, x):
        #ただのインターフェース。
        if chainer.backends.cuda.get_array_module(x) != np:
            #学習の時はここを通る
            return self.static_forward(x)
        else:
            # CPUで実行する場合ここを通る
            return self._forward(x)

のようにしておいて呼び出し側は　model.forward(np.array([x]).astype(np.float32)) のインターフェースを守るという設計にしておけば良いかなと考えています。

まとめ

Static Subgraph Optimizations の機能を触りました。デコレートするだけで簡単に使えます（ここ大事）。
それを実感するサンプルコードを書きました。いろいろ変更すると何がダメでなにがうまく行くかを体感することができます。
この機能がうまく使えるケースがあればあなたは幸運です。十分有効活用してChainerを布教させてください。
V4とV5でコードを両立するテクニックを紹介しました。

おまけ

今回の記事では関係ないですが、
V5 で L.DepthwiseConvolution2D の実装が変わったらしく、GPUでの推論が爆遅になるケースがありますが chainer.config.autotune=True を適用することで改善します。

以上！！！