xchainerをGPU対応させてみた

このたび、弊社ではchainerのScikit-leran向けラッパーであるxchainerのGPU対応, chainer v1.6対応を試みてみました。
chainerのウリでもある複数GPUへの対応はまだまだ道半ばですが、単一GPUへの対応はそれなりに動くようになりましたので、以下に報告します。

この記事は、xchainerのGPU対応に関する技術的な記事です。xchainerを使ってみた記事は別に作成中です。

by Tetsuro Kitajima 2016/2/16

chainer / xchainerとは何か?

少し前にホットな話題になったchainerについてはご存じの方も多いと思います。chainerはPreferred Networks社が開発した深層学習(Deep Learning)向けのフレームワークで、深層学習に用いられるニューラルネットワークを直感的に記述できることで人気を博しました。
chainerや、弊社でも多数の使用実績を持つcaffe, googleが開発したTensorFlowといったフレームワークは実業務の規模で深層学習を取り扱うには不可欠なツールとなっていますが、一方でどのフレームワークも入力・出力、結果の取扱いには独自のノウハウが必要で、使いこなせるようにはそれなりの習得コストを必要とします。
一方、昔ながらの機械学習手法(ランダムツリー、SVM、AdaBoostなど)をpythonで活用しようとする向きにはscikit-learnというフレームワークが公開されています。このモジュールは機械学習手法を問わず統一されたインタフェースを提供しており、ある機械学習手法で一度使い方を習得してしまえばほかの学習手法を使うときにも追加の習得コストが低くなります。
xchainerとは、Scikit-learnとのアダプターを提供し、chainerによる学習プロセスをScikit-learnの評価モジュールを用いて記述できるようにするためのモジュールです。このモジュールを用いることで、学習評価を容易に行えるようになり、Deep Learningと他手法との比較も容易となります。
オリジナルのxchainerはリクルートテクノロジーズ社が公開したものです。

blogエントリ
githubリポジトリ

xchainerは確かにchainerをscikit-learnの枠組みで取り扱える便利なモジュールなのですが、オリジナル版ではGPUを用いた処理には対応しておりません。GPUが使えるようになれば実業務規模のDeep Learningでもxchainerを活用していけるようになるので、今回の試みでGPU対応、ついでに最近のバージョンのchainerへの対応も実装しました。

chainer 1.6対応

最初に、GPU対応の前段階として、最新のchainerでxchainerが動作するか確認します。
まずは普通にchainer1.6環境でxchainerをテストしてみました。

(xchainer-wazalabo)kitajima@machine:~/work/xchainer-wazalabo/xchainer$ python -m unittest discover -s tests
Loading MNIST data for test. This could take a while...
...done

/home/kitajima/work/xchainer-wazalabo/xchainer-wazalabo/local/lib/python2.7/site-packages/chainer/function_set.py:62: FutureWarning: 'collect_parameters' is deprecated. You can pass FunctionSet itself to 'optimizer.setup'
  warnings.warn(msg, FutureWarning)
===Test `fit` method===
This could take a while...
...done

.===Test `predict` method===
This could take a while...
...done

.===Test learning with `cross_val_score` of sklearn===
logging learning process below...
[0 epoch] mean loss: 2.283987, mean accuracy: 0.138827
[1 epoch] mean loss: 2.225600, mean accuracy: 0.258334
[2 epoch] mean loss: 2.159801, mean accuracy: 0.327371
[3 epoch] mean loss: 2.078537, mean accuracy: 0.392072
[4 epoch] mean loss: 1.983404, mean accuracy: 0.491342
[0 epoch] mean loss: 2.362616, mean accuracy: 0.125427
[1 epoch] mean loss: 2.328092, mean accuracy: 0.206697
[2 epoch] mean loss: 2.269016, mean accuracy: 0.336285
[3 epoch] mean loss: 2.182708, mean accuracy: 0.472769
[4 epoch] mean loss: 2.084115, mean accuracy: 0.567456
[ 0.51769848  0.58359288]
...complete

.===Test `execute` method===
F===Test `getFunctions` method===
...done

.===Test `insource` method===
F===Test `outsource` method===
F
======================================================================
FAIL: test_execute (test_packer.NNpackerTestCase)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/kitajima/work/xchainer-wazalabo/xchainer/tests/test_packer.py", line 68, in test_execute
    result = self.nnp.execute(datasets)
  File "xchainer/packer.py", line 28, in execute
    self.entryDatas = self.insource(datasets)
  File "xchainer/packer.py", line 35, in insource
    entries = [{ep: Variable(datasets[ep])} for ep in self.entryPoints]
  File "/home/kitajima/work/xchainer-wazalabo/xchainer-wazalabo/local/lib/python2.7/site-packages/chainer/variable.py", line 48, in __init__
    assert isinstance(data, (numpy.ndarray, cuda.ndarray))
AssertionError

======================================================================
FAIL: test_insource (test_packer.NNpackerTestCase)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/kitajima/work/xchainer-wazalabo/xchainer/tests/test_packer.py", line 50, in test_insource
    entryDatas = self.nnp.insource(datasets)
  File "xchainer/packer.py", line 35, in insource
    entries = [{ep: Variable(datasets[ep])} for ep in self.entryPoints]
  File "/home/kitajima/work/xchainer-wazalabo/xchainer-wazalabo/local/lib/python2.7/site-packages/chainer/variable.py", line 48, in __init__
    assert isinstance(data, (numpy.ndarray, cuda.ndarray))
AssertionError

======================================================================
FAIL: test_outsource (test_packer.NNpackerTestCase)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/kitajima/work/xchainer-wazalabo/xchainer/tests/test_packer.py", line 58, in test_outsource
    childrenOutput = self.nnp.outsource(datasets, False)
  File "xchainer/packer.py", line 40, in outsource
    chldOut = [{chn: ch.execute(datasets, train)} for chn, ch in chld]
  File "xchainer/packer.py", line 28, in execute
    self.entryDatas = self.insource(datasets)
  File "xchainer/packer.py", line 35, in insource
    entries = [{ep: Variable(datasets[ep])} for ep in self.entryPoints]
  File "/home/kitajima/work/xchainer-wazalabo/xchainer-wazalabo/local/lib/python2.7/site-packages/chainer/variable.py", line 48, in __init__
    assert isinstance(data, (numpy.ndarray, cuda.ndarray))
AssertionError

----------------------------------------------------------------------
Ran 7 tests in 14.271s

FAILED (failures=3)
(xchainer-wazalabo)kitajima@machine:~/work/xchainer-wazalabo/xchainer$

以上から分かる通り、警告1件とエラー3件が起きます。まずはここから潰していきましょう。

警告が出ないように chainer1.5以降向けへ改修

警告の原因は、メッセージ通り collect_parameters()を使っていることです。


/home/kitajima/work/xchainer-wazalabo/xchainer-wazalabo/local/lib/python2.7/site-packages/chainer/function_set.py:62: FutureWarning: 'collect_parameters' is deprecated. You can pass FunctionSet itself to 'optimizer.setup

こちらは順当に以下の修正で対応します。

chainer 1.6環境下で警告が出ないように, manager.py を修正

この修正を入れると、警告メッセージは消えます。
ですが、xchainerには他にもchainer1.5以降のAPI変更でdeprecatedになった機能を使っている場所があります。こちらも修正していきましょう。

Chain, Functions -> Linksの変更(chainer1.5対応)を投入
packerの試験コードにもchainer 1.5以後の変更を反映

これで警告は出なくなります。次は本命のエラー対応。

テスト失敗をなくす

先ほどのテスト失敗結果を再掲します。

======================================================================
FAIL: test_outsource (test_packer.NNpackerTestCase)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/kitajima/work/xchainer-wazalabo/xchainer/tests/test_packer.py", line 58, in test_outsource
    childrenOutput = self.nnp.outsource(datasets, False)
  File "xchainer/packer.py", line 40, in outsource
    chldOut = [{chn: ch.execute(datasets, train)} for chn, ch in chld]
  File "xchainer/packer.py", line 28, in execute
    self.entryDatas = self.insource(datasets)
  File "xchainer/packer.py", line 35, in insource
    entries = [{ep: Variable(datasets[ep])} for ep in self.entryPoints]
  File "/home/kitajima/work/xchainer-wazalabo/xchainer-wazalabo/local/lib/python2.7/site-packages/chainer/variable.py", line 48, in __init__
    assert isinstance(data, (numpy.ndarray, cuda.ndarray))
AssertionError

Assertionで文句を言われています。chainerのgithubリポジトリを精査すると、 2015/06/29のコミット Insert type check to Variable.__init__ にて当該のAssertionが加えられたのが分かります。xchainerの元記事が2015/09/02付けなので、おそらくこのコミット前のchainerソースツリーをもとにxchainerを開発していたのだろう、と想像しました。
というわけで、修正していきましょう。かなり荒っぽい修正なのは承知の上で、ざっくりコードに手を入れます。(この修正が適切かどうかはなんとも言えませんが、chainerでなぜこのassertionが入れられたのか不明なので評価が難しいです。chainer側の事情をご存じの方がいらっしゃいましたら、ぜひIssueを立てるか、プルリクをお願いします。)

Chainer側のAssertionに対応するように、テストに供するデータのデータ構造を変更。

試験してみましょう。

(xchainer-wazalabo)kitajima@machine:~/work/xchainer-wazalabo/xchainer$ python -m unittest discover -s tests
Loading MNIST data for test. This could take a while...
...done

===Test `fit` method===
This could take a while...
...done

.===Test `predict` method===
This could take a while...
...done

.===Test learning with `cross_val_score` of sklearn===
logging learning process below...
[0 epoch] mean loss: 2.292086, mean accuracy: 0.141287
[1 epoch] mean loss: 2.262030, mean accuracy: 0.229385
[2 epoch] mean loss: 2.226170, mean accuracy: 0.292936
[3 epoch] mean loss: 2.186005, mean accuracy: 0.330885
[4 epoch] mean loss: 2.139053, mean accuracy: 0.353089
[0 epoch] mean loss: 2.372298, mean accuracy: 0.089392
[1 epoch] mean loss: 2.332475, mean accuracy: 0.216677
[2 epoch] mean loss: 2.272280, mean accuracy: 0.363445
[3 epoch] mean loss: 2.196337, mean accuracy: 0.442191
[4 epoch] mean loss: 2.087648, mean accuracy: 0.531602
[ 0.3569694   0.52675944]
...complete

.===Test `execute` method===
...done

.===Test `getFunctions` method===
...done

.===Test `insource` method===
...done

.===Test `outsource` method===
...done

.
----------------------------------------------------------------------
Ran 7 tests in 14.890s

OK

ひととおり試験が通りました。

GPU対応

ここからは本題のGPU対応化。GPU対応には以下の変更が必要です。

  • テストコードのGPU対応
  • Model, DataのGPU対応

順次手がけていきましょう。

テストコードのGPU対応

まずはGPU対応したコードがどんなインタフェースを持つべきか考えてみます。
現実的に、DeepLearningを仕事で使おうとするとGPU利用がデフォルトになるはずなので、現行のコードそのままでGPU対応ができることが望ましいです。そう考えると、逆にCPU対応ではgpu=Falseといった指定でCPUモードに切り替えられる形でしょうか。

上記の仕様をもとに、テストコードに手を入れてみました。コミットはこちらになります。

GPU対応のテストコードをコミット

なお、上記のテストコードだと本来は不十分です:GPU対応できていない状態のソースツリーでもテストに通ってしまいます。
本来は大規模なDeep Learning問題に対してGPU対応で高速化できていることも検出すべきなのですが、今回はGPU対応コードでも動いていることを立証する意図でのテストとして書きました。

Model, DataのGPU対応

テストコードだけGPU対応しても意味がありません。というわけで、Model, Dataを取り扱う部分をGPU対応します。

CPU版とGPU版でのDL学習・判定処理の流れは、大きく見て以下のようになっています。

今回は上の図で, GPUメモリへのコピー、メインメモリへのコピー部分を追加する形になります。幸いchainerにはこれらの処理のためのAPIがあるので、実装は容易です。

対応したコードはこちらにコミットしてあります。

ManagerのGPU対応
テストを通じて, GPU利用でパフォーマンスが上がっていることを確認しましょう。

python -m unittest discover -s tests -p 'test_manager.py'
(中略)
Ran 3 tests in 14.367s

python -m unittest discover -s tests -p 'test_gpu_manager.py'
(中略)
Ran 3 tests in 2.434s

これで、学習はGPUを使って進められるようになりました。

実際に使ってみた

ここまでの改修を入れたバージョンのxchainer、しばらく実験で使っていましたが、いくつか不便な点がありました。

  • 学習を進めている最中に、学習データのloss, accuracyだけではなく、テストデータのloss, accuracyも見たい。テストデータのloss, accuracyを見られないと、過学習の判定ができない。
  • ログのフォーマットも変更できるようにしておきたい。

このあたりも、改修してみました。もちろん、改修に伴いドキュメントの修正が必要ですのでこちらも手を入れて。
ついでに, chainer r1.6では動かなくなっていたサンプルコード mnist_complex.py も削除しました。

一通りの改修を加えた結果が、こちらになります。

便利系の機能(ログファイルのフォーマット変更, 学習中のテストデータによる評価結果表示)を追加

最新版の出力結果の例を以下に示します。5 epochsごとにTesting loss/accuracyを表示するように設定しました。


[0 epoch] mean loss: 0.969828, mean accuracy: 0.603322, testing loss: 0.931455, testing accuracy: 0.631683
[1 epoch] mean loss: 0.912616, mean accuracy: 0.631805
[2 epoch] mean loss: 0.882430, mean accuracy: 0.643660
[3 epoch] mean loss: 0.849373, mean accuracy: 0.658796
[4 epoch] mean loss: 0.822387, mean accuracy: 0.672851
[5 epoch] mean loss: 0.796034, mean accuracy: 0.684222, testing loss: 0.796960, testing accuracy: 0.685547
[6 epoch] mean loss: 0.766287, mean accuracy: 0.695962
[7 epoch] mean loss: 0.738852, mean accuracy: 0.704800
[8 epoch] mean loss: 0.712139, mean accuracy: 0.713099
[9 epoch] mean loss: 0.687893, mean accuracy: 0.719742

現在、今回改修を加えたxchainerを用いて機械学習の実験を進めております。こちらは別の記事でご紹介予定。