Chainerの抽象化に挑戦してみた - のんびりしているエンジニアの日記

皆さんこんにちは
お元気ですか？年末ってこんなに忙しくなるものですね。

本記事はChainer Advent Calendar 16日目の記事です。

本日はScikit-learn likeなChainerを作った記事です。なぜ、作ったのかは後述します。

はじめに
- なぜ、抽象化コードを作ってみたのか
- Scikit-learnのWrapperを実装するメリット
何ができれば嬉しいか
- 最低限必要な項目
実際に作ってみた
- ソースコード（chainer_neural_network.py）
MNISTで遊ぶ
- ソースコード
最後に

はじめに

なぜ、抽象化コードを作ってみたのか

Chainerの実装を高速にし、検証の立ち上がりを早くしたいからです。

Chainerは非常にFlexibleなネットワークを作れるので、個人的に好きです。
特にRecurrent Neural Network系のネットワークを書く時に重宝しています。
また、実装はTrainerを使うとTrainer実装前よりも簡単な実装になりました。

しかし、手軽さと最初の立ち上がりの早さが、Scikit-learnと比べるとまだ劣ると感じています。

せっかくなので、アドベントカレンダーを機に抽象化をある程度頑張ってみようかなと思った次第です。

Scikit-learnのWrapperを実装するメリット

Wrapper実装メリットは3つあると考えています。

「うまくいかないこと」を早く知る
使いやすい
他のアルゴリズムと取り替えやすい

「うまくいかないこと」を早く知る

私が考える最も大きな理由は「うまくいかないこと」を早く知れることです。

正直、実装初回からうまくいくとは思っていません。
そのため、早めになぜ、うまくいかないのか、
更には「うまくいかない」ことがわかりたいです。

これを早く知れると次はどうしようかといった手を打ちやすいです。

使いやすい

scikit-learnのインターフェースはシンプルで使いやすいです。
アルゴリズムを使う際にfit, predict, predict_probaあたりが何をしているのかがわかれば十分です。
プログラマはそのメソッドを呼ぶだけで、内部動作を知らなくても期待する結果を得られます。

他のアルゴリズムと取り替えやすい

機械学習アルゴリズムの入れ替えが容易です。
インターフェースを共通化した場合、基本的に初期化が変わるのみで
コードを殆ど変更する必要がありません。
そのため、入れ替えも非常に簡単で、検証する速度が効率化します。

何ができれば嬉しいか

必要な作業の洗い出しをまずやります。
Chainerの機能をフルに活用するのはRecurrent Neural Network込みにすると難しいと思っています。
そのため、一部の機能に絞れば、きっと実装できるはず・・・
ということで必要最低限の機能を真面目に考えてみました。

最低限必要な項目

モデル定義・・・ニューラルネットワーク定義を入力する。
ハイパーパラメータ・・・学習用パラメータ
GPU or CPU・・・ハードのモードの切り替え
学習機能・・・モデル学習機能（fit）
出力・・・確率分布での出力（predict_proba）、ラベル出力（predict）
save load機能(weight)・・・学習済みモデルの読み出し、書き出し

実際に作ってみた

こんな感じになりました。
可能な限り、既存のChainerの機能を活かすようにしています。

ソースコード（chainer_neural_network.py）

# coding:utf-8
from __future__ import absolute_import
from __future__ import unicode_literals
from sklearn.base import BaseEstimator
import chainer
import numpy as np
from chainer.training import extensions
import cupy


class ChainerNeuralNetwork(BaseEstimator):
    """
    Chainer Neural Network
    Scikit learn Wrapper
    """

    def __init__(self, model, optimizer, extensions=[], gpu=-1, batch_size=16, epochs=10, result="result"):
        self.model = model
        self.optimizer = optimizer
        self.extensions = extensions
        self.gpu = gpu
        self.batch_size = batch_size
        self.epochs = epochs
        self.result = result

        self.xp = np
        if gpu >= 0:
            chainer.cuda.get_device(gpu).use()
            self.model.to_gpu()
            self.xp = cupy
        self.optimizer.setup(model)
        super(ChainerNeuralNetwork, self).__init__()

    def convert_list_to_tuple(self, X, y):
        """
        convert list to tuple

        :param X: X
        :param y: y
        :return: tuple dataset
        """
        return chainer.datasets.TupleDataset(X, y)

    def fit(self, X, y=None, valid_X=None, valid_y=None):
        """
        fit model

        :param X: tuple dataset or matrix X
        :param y: expected vector
        :param valid_X: validation X
        :param valid_y: validation y
        :return:
        """
        if isinstance(X, chainer.datasets.TupleDataset):
            train_dataset = X
        else:
            train_dataset = self.convert_list_to_tuple(X, y)
        train_iter = chainer.iterators.SerialIterator(train_dataset, self.batch_size)

        updater = chainer.training.StandardUpdater(train_iter, self.optimizer, device=self.gpu)
        trainer = chainer.training.Trainer(updater, (self.epochs, 'epoch'), out=self.result)

        if valid_X is not None:
            if isinstance(valid_X, chainer.datasets.TupleDataset):
                valid_dataset = valid_X
            else:
                valid_dataset = self.convert_list_to_tuple(valid_X, valid_y)
            test_iter = chainer.iterators.SerialIterator(valid_dataset, self.batch_size,
                                                         repeat=False, shuffle=False)
            self.extensions.append(extensions.Evaluator(test_iter, self.model, device=self.gpu))

        for extend in self.extensions:
            trainer.extend(extend)
        trainer.run()

    def _iter_predict(self, X):
        """
        iteration for prediction

        :param X: np.ndarray
        :return: prediction of batch
        """
        for index in range(0, len(X), self.batch_size):
            batch_x = chainer.Variable(self.xp.array(X[index: index + self.batch_size]).astype(self.xp.float32),
                                       volatile=False)
            if index == 0:
                predicts = np.array(chainer.cuda.to_cpu(self.model.predict(batch_x).data))
            else:
                predicts = np.vstack([predicts, chainer.cuda.to_cpu(self.model.predict(batch_x).data)])
        return np.array(predicts)

    def predict(self, X):
        return np.argmax(self.predict_proba(X), axis=1)

    def predict_proba(self, X):
        return self._iter_predict(X)

    def load_weights(self, filepath):
        chainer.serializers.load_npz(filepath, self.model)

    def save_weights(self, filepath):
        chainer.serializers.save_npz(filepath, self.model)

基本的にはfit, predict, predict_probaの基本機能を実装しました。
学習して保存できるようload_weights, save_weights機能を作ってみました。
X, yを別々にした場合でも、Dataset（Tuple）でも使えるようにしています。

MNISTで遊ぶ

Wrapperを使ってチュートリアルのMNISTコードを抽象化します。

簡単なニューラルネットワークであればこの程度抽象化すれば実施可能です。
ただし、構築モデル（Chain）はモデルを使うためにpredict関数の実装をしなければなりません。
そのため、なければ実装してください。

本プログラムは、chainer.links.Classifierを継承し、ClassifierWrapperを作りました。
このクラスにpredict関数を実装しています。

ソースコード

# coding:utf-8
from __future__ import absolute_import
from __future__ import unicode_literals
import argparse
import chainer
import chainer.functions as F
import chainer.links as L
from chainer import training
from chainer.training import extensions
from chainer_neural_network import ChainerNeuralNetwork
import chainer.serializers
import numpy as np


class MLP(chainer.Chain):
    def __init__(self, n_units, n_out):
        super(MLP, self).__init__(
            l1=L.Linear(None, n_units),  
            l2=L.Linear(None, n_units),  
            l3=L.Linear(None, n_out),
        )

    def __call__(self, x):
        h1 = F.relu(self.l1(x))
        h2 = F.relu(self.l2(h1))
        return self.l3(h2)


class ClassifierWrapper(L.Classifier):
    def predict(self, x):
        return F.softmax(self.predictor(x))


def main():
    parser = argparse.ArgumentParser(description='Chainer example: MNIST')
    parser.add_argument('--batchsize', '-b', type=int, default=100,
                        help='Number of images in each mini-batch')
    parser.add_argument('--epoch', '-e', type=int, default=20,
                        help='Number of sweeps over the dataset to train')
    parser.add_argument('--gpu', '-g', type=int, default=-1,
                        help='GPU ID (negative value indicates CPU)')
    parser.add_argument('--out', '-o', default='result',
                        help='Directory to output the result')
    parser.add_argument('--unit', '-u', type=int, default=1000,
                        help='Number of units')
    args = parser.parse_args()

    print('GPU: {}'.format(args.gpu))
    print('# unit: {}'.format(args.unit))
    print('# Minibatch-size: {}'.format(args.batchsize))
    print('# epoch: {}'.format(args.epoch))
    print('')

    model = ClassifierWrapper(MLP(args.unit, 10))
    optimizer = chainer.optimizers.Adam()

    extensions_list = [
        extensions.PrintReport(
            ['epoch', 'main/loss', 'validation/main/loss',
             'main/accuracy', 'validation/main/accuracy']),
        extensions.LogReport(),
        extensions.dump_graph('main/loss'),
        extensions.ProgressBar()
    ]

    train, test = chainer.datasets.get_mnist()
    X, y = [pair[0] for pair in train], [pair[1] for pair in train]

    test_X, tedt_y = [pair[0] for pair in test], [pair[1] for pair in test]
    test_X, tedt_y = np.array(test_X), np.array(tedt_y)

    neural_network = ChainerNeuralNetwork(model=model, optimizer=optimizer, extensions=extensions_list, gpu=args.gpu,
                                          batch_size=args.batchsize, epochs=args.epoch, result=args.out)
    neural_network.fit(X=X, y=y, valid_X=test)
    print(neural_network.predict_proba(test_X).shape)
    print(neural_network.predict(test_X).shape)

    neural_network.save_weights("model.npz")
    neural_network.load_weights("model.npz")

if __name__ == '__main__':
    main()