手のひらスパコンでKeras-Tensorflowを使ってみる

この記事は Retty Advent Calendar 20日目です。
昨日は@takumi-suzukiのDNSでバランシングしていたら辛くなった話でした。

#ガジェット手当
みなさん、いきなりですがスパコンを持ってますか？
僕は持っています。

Rettyという会社は、エンジニアには通期で一回「ガジェット手当」という制度があります。
IT機器などを個人的に購入する場合に、業務とか用途関係なく援助してくれる制度です。

だからというわけではありませんが、Rettyでもスパコンを買ってみることにしました。

#スパコン買ってみよう！

皆さんが想像されるスパコンというのはこういうのをイメージされると思いますが、

By 0-0t - 0-0t, GFDL-no-disclaimers, Link

やっぱ買うならこういうのが欲しいなと思ったりするのですが、残念ながらまだまだベンチャーのRettyにはこんな広いスペースは確保できません。
僕の個人的な趣味で、みんなの快適なオフィスを占有するのは少々心が痛みます。

実はこの世界の進歩はものすごく早くて、最近では株式会社PEZY Computingという会社は、２０２０年までに京の２台分の演算能力を持ったスパコンをコピー機サイズにして、オフィスで一人一台が使えるようにするということを表明していて話題になったりしています。

とはいえ、月額の電気代が１５万もするコピー機は、毎月のAWSの請求にドキドキしているような私にはどう考えても扱えそうにありません。。。

そもそもスーパーコンピューターって何って考えたところ

一般的に使用されるサーバ機よりも浮動小数点演算が1,000倍以上の速さのコンピュータを「スーパーコンピュータ」と呼ぶことが多い。

らしいです。
浮動小数点演算といえば、CUDAに代表されるGPGPUが有名ですね。
と言うので、NVIDIAのサイトを見ていたらこんなものを見つけました。

NVIDIA Jetson TK1

NVIDIA Jetson TK1
ちっちゃいですね！

世界中のスーパーコンピュータに組み込まれているのと同じ NVIDIA Kepler™ アーキテクチャのコンピューティングコアを使用しています。

もともと組み込みの機械学習システム開発の検証用ボードみたいなんですが、これならみんなに迷惑をかけずに憧れのスパコンオーナーになれそうです。

CUDA環境のセットアップ

諸々の手順をかなり端折りますが、JetsonにJetPack L4TからCUDA環境をセットアップします。

おお、CUDAドライバーがインストールされています。

ubuntu@tegra-ubuntu:~$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2014 NVIDIA Corporation
Built on Tue_Feb_17_22:53:16_CST_2015
Cuda compilation tools, release 6.5, V6.5.45

##とりあえず機械学習でもやってみるか・・・
巷では最近、機械学習やらAIやらが流行りらしいですが、Rettyでは普通にいろんなところで機械学習で得られたモデルをサービスに活用しています。
当社の機械学習エンジニアに聞くと、最近はKerasというのが流行っているらしいのでそれを試して見ます。

###Tensorflowのセットアップ
CUDAが利用可能なので、ソースコードからビルドします。
基本はここに書いてあるとうりですが、今回はVirtualenvでインストールします。

$ sudo apt-get install python-pip python-dev python-virtualenv python-gtk2-dev gfortran swig liblapack-dev

virtualenvの環境を作って、アクティベート。

$ virtualenv --system-site-packages ~/tensorflow
$ source ~/tensorflow/bin/activate

今回はCUDAが利用可能なのですが、最新のバイナリはCUDA-8.0、CuDNN-v5が必要です。

# Ubuntu/Linux 64-bit, GPU enabled, Python 2.7
# Requires CUDA toolkit 8.0 and CuDNN v5. For other versions, see "Installing from sources" below.

TK1のCUDA toolkitは6.5、CuDNNはv2になっているのとCPUがARMコアなので、ソースコードからコンパイルします。
その前にbazelのインストールをしておきます。

$ git clone --recurse-submodules https://github.com/tensorflow/tensorflow
$ cd tensorflow/tensorflow
$ ./configure
Do you wish to bulid TensorFlow with GPU support? [y/n] y
GPU support will be enabled for TensorFlow

Please specify the location where CUDA 6.5 toolkit is installed. Refer to
README.md for more details. [default is: /usr/local/cuda]: /usr/local/cuda-6.5

Please specify the location where CUDNN 6.5 V2 library is installed. Refer to
README.md for more details. [default is: /usr/local/cuda]: /usr/local/cuda-6.5

Setting up Cuda include
Setting up Cuda lib64
Setting up Cuda bin
Setting up Cuda nvvm
Configuration finished

###kerasのセットアップ
Keras: Deep Learning library for Theano and TensorFlow

$ sudo pip install keras

これだけです。

###MNISTで動作確認

Kerasの開発者François CholletさんのGithubからサンプルコードをダウンロードして実行して見ます。

mnist_cnn.py

'''Trains a simple convnet on the MNIST dataset.
Gets to 99.25% test accuracy after 12 epochs
(there is still a lot of margin for parameter tuning).
16 seconds per epoch on a GRID K520 GPU.
'''

from __future__ import print_function
import numpy as np
np.random.seed(1337)  # for reproducibility

from keras.datasets import mnist
from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation, Flatten
from keras.layers import Convolution2D, MaxPooling2D
from keras.utils import np_utils
from keras import backend as K

batch_size = 128
nb_classes = 10
nb_epoch = 12

# input image dimensions
img_rows, img_cols = 28, 28
# number of convolutional filters to use
nb_filters = 32
# size of pooling area for max pooling
pool_size = (2, 2)
# convolution kernel size
kernel_size = (3, 3)

# the data, shuffled and split between train and test sets
(X_train, y_train), (X_test, y_test) = mnist.load_data()

if K.image_dim_ordering() == 'th':
    X_train = X_train.reshape(X_train.shape[0], 1, img_rows, img_cols)
    X_test = X_test.reshape(X_test.shape[0], 1, img_rows, img_cols)
    input_shape = (1, img_rows, img_cols)
else:
    X_train = X_train.reshape(X_train.shape[0], img_rows, img_cols, 1)
    X_test = X_test.reshape(X_test.shape[0], img_rows, img_cols, 1)
    input_shape = (img_rows, img_cols, 1)

X_train = X_train.astype('float32')
X_test = X_test.astype('float32')
X_train /= 255
X_test /= 255
print('X_train shape:', X_train.shape)
print(X_train.shape[0], 'train samples')
print(X_test.shape[0], 'test samples')

# convert class vectors to binary class matrices
Y_train = np_utils.to_categorical(y_train, nb_classes)
Y_test = np_utils.to_categorical(y_test, nb_classes)

model = Sequential()

model.add(Convolution2D(nb_filters, kernel_size[0], kernel_size[1],
                        border_mode='valid',
                        input_shape=input_shape))
model.add(Activation('relu'))
model.add(Convolution2D(nb_filters, kernel_size[0], kernel_size[1]))
model.add(Activation('relu'))
model.add(MaxPooling2D(pool_size=pool_size))
model.add(Dropout(0.25))

model.add(Flatten())
model.add(Dense(128))
model.add(Activation('relu'))
model.add(Dropout(0.5))
model.add(Dense(nb_classes))
model.add(Activation('softmax'))

model.compile(loss='categorical_crossentropy',
              optimizer='adadelta',
              metrics=['accuracy'])

model.fit(X_train, Y_train, batch_size=batch_size, nb_epoch=nb_epoch,
          verbose=1, validation_data=(X_test, Y_test))
score = model.evaluate(X_test, Y_test, verbose=0)
print('Test score:', score[0])
print('Test accuracy:', score[1])

おお動いた！
ちょっと遅いけど・・・

$ python minst_cnn.py 
Using TensorFlow backend.
X_train shape: (60000, 28, 28, 1)
60000 train samples
10000 test samples
Train on 60000 samples, validate on 10000 samples
Epoch 1/12
60000/60000 [==============================] - 147s - loss: 0.3819 - acc: 0.8820 - val_loss: 0.0856 - val_acc: 0.9740
Epoch 2/12
60000/60000 [==============================] - 154s - loss: 0.1335 - acc: 0.9605 - val_loss: 0.0617 - val_acc: 0.9794
Epoch 3/12
60000/60000 [==============================] - 150s - loss: 0.1042 - acc: 0.9689 - val_loss: 0.0529 - val_acc: 0.9833
Epoch 4/12
60000/60000 [==============================] - 149s - loss: 0.0863 - acc: 0.9737 - val_loss: 0.0446 - val_acc: 0.9852
Epoch 5/12
60000/60000 [==============================] - 147s - loss: 0.0785 - acc: 0.9768 - val_loss: 0.0402 - val_acc: 0.9864
Epoch 6/12
60000/60000 [==============================] - 155s - loss: 0.0703 - acc: 0.9790 - val_loss: 0.0387 - val_acc: 0.9871
Epoch 7/12
60000/60000 [==============================] - 157s - loss: 0.0627 - acc: 0.9817 - val_loss: 0.0371 - val_acc: 0.9875
Epoch 8/12
60000/60000 [==============================] - 157s - loss: 0.0602 - acc: 0.9820 - val_loss: 0.0353 - val_acc: 0.9885
Epoch 9/12
60000/60000 [==============================] - 156s - loss: 0.0557 - acc: 0.9833 - val_loss: 0.0341 - val_acc: 0.9886
Epoch 10/12
60000/60000 [==============================] - 155s - loss: 0.0539 - acc: 0.9838 - val_loss: 0.0320 - val_acc: 0.9892
Epoch 11/12
60000/60000 [==============================] - 145s - loss: 0.0510 - acc: 0.9850 - val_loss: 0.0307 - val_acc: 0.9898
Epoch 12/12
60000/60000 [==============================] - 165s - loss: 0.0459 - acc: 0.9862 - val_loss: 0.0318 - val_acc: 0.9897
Test score: 0.0317565297928
Test accuracy: 0.9897  

今回はCUDA無しの環境でテストしていないですが、年末にでも消費電力や処理時間などベンチマークを取ってみたいと思います。
また、Jetson TK1の後継機でプロダクションモデルのJetson TX1というのがすでに製品化されていて、実際に自動運転車やドローンなどに使われているようです。
次回の「ガジェット手当」が楽しみです。

##終わりに
ちなみに、RettyではAKIBAという機械学習の為のCUDA環境を自前で用意していて、エンジニアに限らず誰もが利用することができます。
また、エンジニアの可能性を無限大に伸ばす為に、新しい働き方や価値観を支援する様々な制度を準備しています。
もし興味があれば遊びに来てください！