Chainer:GPUを使うプログラムはどう書くか

ChainerでGPUを使ってみようで、MNISTのデータセットを使ったディープラーニングのサンプルプログラムである train_mnist.py のコマンドで、-gpu 0 でGPUが使えることを示した。
つまり、操作説明をしたに過ぎなかったが、今回は、プログラムレベルでどう対応しているか見てみよう。
プログラムは,https://github.com/pfnet/chainer/tree/master/examples/mnist である。

    parser.add_argument('--gpu', '-g', type=int, default=-1,
help='GPU ID (negative value indicates CPU)')

ここで、 -gpu の処理が行われ、args.gpu に整数値が入る。
デフォルト値は -1 で、GPUを使用しないことを示す。

    if args.gpu >= 0:
chainer.cuda.get_device(args.gpu).use()  # Make a specified GPU current
model.to_gpu()  # Copy the model to the GPU

GPUを使用するときだけ、指定したGPUを確保し、作ったモデルをGPUにコピーするようだ。

    # Set up a trainer
updater = training.StandardUpdater(train_iter, optimizer, device=args.gpu)
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)

# Evaluate the model with the test dataset for each epoch
trainer.extend(extensions.Evaluator(test_iter, model, device=args.gpu))

あとは、上記の２個所で、device指定があるので、args.gpu をそのまま与えている。
GPUのの使用に関係してくるのは、たったこれだけである。
 あとは、chainerがよきに計らってくれるので、何もすることはない。
フレームワークとは、こういうものである。

ということで、よくわからなくてもGPUを使うプログラムが書けてしまう。
便利というか、便利過ぎな気がする。

そのうち、Python から直接GPUいじるためにCuPyを紹介するかもしれない。 NumPy互換のインターフェイスでCUDA（NVIDIAのGPUライブラリ）を使えるらしい。