WonderPlanet Tech Blog

アドバンストテクノロジー部やCTO室のメンバーを中心に最新の技術情報を発信しています。

Kerasで複数のGPUを使い、学習を高速化してみた

こんにちは。アドバンストテクノロジー部のR&Dチーム所属岩原です。
今回はKerasで複数のGPUを使う方法を書きたいと思います。

Keras 2.0.9から簡単に複数GPUを使用した高速化が可能に。

Keras2.0.9からtraining_utilsというモジュールにmulti_gpu_modelという関数が追加されました。
コレを使うと、学習を複数のGPUで行わせることが可能になります。
inputをGPUの数だけ分割することによって、並列化を実現しているようです。
keras/training_utils.py at master · keras-team/keras

では、実際に試してみましょう。

環境

  • AWS EC2(p2.8xlarge) -> GPU8本
  • Deep Learning Base AMI (Ubuntu) Version 2.0 (ami-041db87c) -> CUDAやCuDNN、Python3.5などがセットアップ済み
  • Keras 2.1.2
  • tensorflow-gpu 1.4.1
  • Python 3.5.2

今回使用するコード

Kerasの例にあるcifar10_cnn.pyを複数GPUに対応させてみたいと思います。
keras/cifar10_cnn.py at master · keras-team/keras

まずはGPU1つのみの場合はどれくらいかかったのかを以下に示します。

~~~ 省略 ~~~
Epoch 90/100
1563/1563 [==============================] - 21s 13ms/step - loss: 0.7667 - acc: 0.7457 - val_loss: 0.6644 - val_acc: 0.7838
Epoch 91/100
1563/1563 [==============================] - 21s 13ms/step - loss: 0.7694 - acc: 0.7457 - val_loss: 0.6627 - val_acc: 0.7803
Epoch 92/100
1563/1563 [==============================] - 21s 13ms/step - loss: 0.7692 - acc: 0.7449 - val_loss: 0.7553 - val_acc: 0.7680
Epoch 93/100
1563/1563 [==============================] - 21s 13ms/step - loss: 0.7721 - acc: 0.7448 - val_loss: 0.7210 - val_acc: 0.7862
Epoch 94/100
1563/1563 [==============================] - 21s 13ms/step - loss: 0.7751 - acc: 0.7436 - val_loss: 0.6743 - val_acc: 0.7811
Epoch 95/100
1563/1563 [==============================] - 21s 13ms/step - loss: 0.7781 - acc: 0.7412 - val_loss: 0.7047 - val_acc: 0.7725
Epoch 96/100
1563/1563 [==============================] - 21s 13ms/step - loss: 0.7781 - acc: 0.7427 - val_loss: 0.6371 - val_acc: 0.7909
Epoch 97/100
1563/1563 [==============================] - 21s 13ms/step - loss: 0.7720 - acc: 0.7452 - val_loss: 0.6331 - val_acc: 0.7949
Epoch 98/100
1563/1563 [==============================] - 21s 13ms/step - loss: 0.7917 - acc: 0.7399 - val_loss: 0.7105 - val_acc: 0.7699
Epoch 99/100
1563/1563 [==============================] - 21s 13ms/step - loss: 0.7829 - acc: 0.7414 - val_loss: 0.6481 - val_acc: 0.7859
Epoch 100/100
1563/1563 [==============================] - 21s 13ms/step - loss: 0.7868 - acc: 0.7404 - val_loss: 0.6266 - val_acc: 0.8005

ということで、学習に21 sec / 1epoch かかっていることになります。
では、コレを単純に8つのGPUで実行した場合にどのような結果になるのか試してみましょう。

単純に8つのGPUで実行した場合

importの追加

まずは必要なモジュールのimportを追加します。
モデルのビルドはCPUで行う必要があるため、Tensorflowをimportします。

import tensorflow as tf # add
from keras.utils.training_utils import multi_gpu_model # add

GPU数の定数を追加とbatch_sizeの変更

GPUの数を定数として定義しておきます。今回は8つ使用するので、8を指定します。
また、batch_sizeは並列で処理を行うために元々のbatch_sizeをGPUの数だけ掛けます。

gpu_count = 8 # add

batch_size = 32 * gpu_count # modify

モデルの構築はCPUで行う

モデルの構築はOOMエラー対策のため、CPUで明示的に行う必要があるので、
tf.deviceを使用します。

with tf.device("/cpu:0"): # add
    model = Sequential()
    model.add(Conv2D(32, (3, 3), padding='same',
                     input_shape=x_train.shape[1:]))
# 以下略

複数GPU対応する

modelをコンパイルする直前に、multi_gpu_model関数を呼び出すようにします。
引数gpusにはGPUの数を指定します。1を指定すると実行時エラーになるので注意して下さい。

model = multi_gpu_model(model, gpus=gpu_count) # add

コード全体

追記変更を含めたコード全体は以下のとおりです。

'''Train a simple deep CNN on the CIFAR10 small images dataset.

It gets to 75% validation accuracy in 25 epochs, and 79% after 50 epochs.
(it's still underfitting at that point, though).
'''

from __future__ import print_function
import keras
from keras.datasets import cifar10
from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation, Flatten
from keras.layers import Conv2D, MaxPooling2D
import os
import tensorflow as tf # add
from keras.utils.training_utils import multi_gpu_model # add

gpu_count = 8 # add

batch_size = 32 * gpu_count # modify
num_classes = 10
epochs = 100
data_augmentation = True
num_predictions = 20
save_dir = os.path.join(os.getcwd(), 'saved_models')
model_name = 'keras_cifar10_trained_model.h5'


# The data, shuffled and split between train and test sets:
(x_train, y_train), (x_test, y_test) = cifar10.load_data()
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 = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)
with tf.device("/cpu:0"): # add
    model = Sequential()
    model.add(Conv2D(32, (3, 3), padding='same',
                     input_shape=x_train.shape[1:]))
    model.add(Activation('relu'))
    model.add(Conv2D(32, (3, 3)))
    model.add(Activation('relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.25))

    model.add(Conv2D(64, (3, 3), padding='same'))
    model.add(Activation('relu'))
    model.add(Conv2D(64, (3, 3)))
    model.add(Activation('relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.25))

    model.add(Flatten())
    model.add(Dense(512))
    model.add(Activation('relu'))
    model.add(Dropout(0.5))
    model.add(Dense(num_classes))
    model.add(Activation('softmax'))
model = multi_gpu_model(model, gpus=gpu_count) # add
# initiate RMSprop optimizer
opt = keras.optimizers.rmsprop(lr=0.0001, decay=1e-6)

# Let's train the model using RMSprop
model.compile(loss='categorical_crossentropy',
              optimizer=opt,
              metrics=['accuracy'])

x_train = x_train.astype('float32')
x_test = x_test.astype('float32')
x_train /= 255
x_test /= 255

if not data_augmentation:
    print('Not using data augmentation.')
    model.fit(x_train, y_train,
              batch_size=batch_size,
              epochs=epochs,
              validation_data=(x_test, y_test),
              shuffle=True)
else:
    print('Using real-time data augmentation.')
    # This will do preprocessing and realtime data augmentation:
    datagen = ImageDataGenerator(
        featurewise_center=False,  # set input mean to 0 over the dataset
        samplewise_center=False,  # set each sample mean to 0
        featurewise_std_normalization=False,  # divide inputs by std of the dataset
        samplewise_std_normalization=False,  # divide each input by its std
        zca_whitening=False,  # apply ZCA whitening
        rotation_range=0,  # randomly rotate images in the range (degrees, 0 to 180)
        width_shift_range=0.1,  # randomly shift images horizontally (fraction of total width)
        height_shift_range=0.1,  # randomly shift images vertically (fraction of total height)
        horizontal_flip=True,  # randomly flip images
        vertical_flip=False)  # randomly flip images

    # Compute quantities required for feature-wise normalization
    # (std, mean, and principal components if ZCA whitening is applied).
    datagen.fit(x_train)

    # Fit the model on the batches generated by datagen.flow().
    model.fit_generator(datagen.flow(x_train, y_train,
                                     batch_size=batch_size),
                        epochs=epochs,
                        validation_data=(x_test, y_test),
                        workers=4)

# Save model and weights
if not os.path.isdir(save_dir):
    os.makedirs(save_dir)
model_path = os.path.join(save_dir, model_name)
model.save(model_path)
print('Saved trained model at %s ' % model_path)

# Score trained model.
scores = model.evaluate(x_test, y_test, verbose=1)
print('Test loss:', scores[0])
print('Test accuracy:', scores[1])

結果

~~~ 省略 ~~~

Epoch 90/100
196/196 [==============================] - 18s 90ms/step - loss: 0.7307 - acc: 0.7453 - val_loss: 0.6331 - val_acc: 0.7831
Epoch 91/100
196/196 [==============================] - 18s 90ms/step - loss: 0.7301 - acc: 0.7457 - val_loss: 0.6280 - val_acc: 0.7817
Epoch 92/100
196/196 [==============================] - 18s 89ms/step - loss: 0.7324 - acc: 0.7445 - val_loss: 0.6149 - val_acc: 0.7870
Epoch 93/100
196/196 [==============================] - 18s 91ms/step - loss: 0.7239 - acc: 0.7494 - val_loss: 0.6257 - val_acc: 0.7822
Epoch 94/100
196/196 [==============================] - 18s 90ms/step - loss: 0.7247 - acc: 0.7472 - val_loss: 0.6101 - val_acc: 0.7879
Epoch 95/100
196/196 [==============================] - 18s 89ms/step - loss: 0.7207 - acc: 0.7482 - val_loss: 0.6233 - val_acc: 0.7860
Epoch 96/100
196/196 [==============================] - 18s 89ms/step - loss: 0.7191 - acc: 0.7489 - val_loss: 0.6349 - val_acc: 0.7798
Epoch 97/100
196/196 [==============================] - 18s 90ms/step - loss: 0.7111 - acc: 0.7514 - val_loss: 0.6057 - val_acc: 0.7912
Epoch 98/100
196/196 [==============================] - 18s 90ms/step - loss: 0.7118 - acc: 0.7524 - val_loss: 0.6084 - val_acc: 0.7894
Epoch 99/100
196/196 [==============================] - 18s 90ms/step - loss: 0.7126 - acc: 0.7523 - val_loss: 0.6026 - val_acc: 0.7887
Epoch 100/100
196/196 [==============================] - 18s 89ms/step - loss: 0.7035 - acc: 0.7567 - val_loss: 0.6052 - val_acc: 0.7932

21秒よりは少し早くなりましたが、思ってたよりは早くなっていません。

仮説

実際に学習をしている様子を眺めているとわかるのですが、進捗が頻繁に止まります。
これはImageDataGeneratorが画像を作る処理がボトルネックになっていそうです。
画像を生成する処理が追いついていないのではないか、という仮説に基づいて、対応してみましょう。
例のコードではfit_generatorを使うようになっているので、それをfit関数を使ってImageDataGeneratorが関与しないように変更してみましょう。

fit_generatorではなくfitを使うようにした場合

data_augmentation = True

data_augmentation = False

に変えるだけでOKです。

なお、サンプル数は変わらず50000になります。

Epoch 90/100
50000/50000 [==============================] - 6s 113us/step - loss: 0.5075 - acc: 0.8223 - val_loss: 0.6516 - val_acc: 0.7786
Epoch 91/100
50000/50000 [==============================] - 6s 113us/step - loss: 0.4984 - acc: 0.8247 - val_loss: 0.6293 - val_acc: 0.7859
Epoch 92/100
50000/50000 [==============================] - 6s 113us/step - loss: 0.4987 - acc: 0.8237 - val_loss: 0.6290 - val_acc: 0.7860
Epoch 93/100
50000/50000 [==============================] - 6s 113us/step - loss: 0.4925 - acc: 0.8277 - val_loss: 0.6383 - val_acc: 0.7849
Epoch 94/100
50000/50000 [==============================] - 6s 113us/step - loss: 0.4922 - acc: 0.8274 - val_loss: 0.6283 - val_acc: 0.7839
Epoch 95/100
50000/50000 [==============================] - 6s 113us/step - loss: 0.4892 - acc: 0.8292 - val_loss: 0.6435 - val_acc: 0.7832
Epoch 96/100
50000/50000 [==============================] - 6s 113us/step - loss: 0.4850 - acc: 0.8298 - val_loss: 0.6449 - val_acc: 0.7820
Epoch 97/100
50000/50000 [==============================] - 6s 113us/step - loss: 0.4823 - acc: 0.8325 - val_loss: 0.6250 - val_acc: 0.7878
Epoch 98/100
50000/50000 [==============================] - 6s 113us/step - loss: 0.4741 - acc: 0.8337 - val_loss: 0.6227 - val_acc: 0.7902
Epoch 99/100
50000/50000 [==============================] - 6s 113us/step - loss: 0.4692 - acc: 0.8345 - val_loss: 0.6559 - val_acc: 0.7794
Epoch 100/100
50000/50000 [==============================] - 6s 113us/step - loss: 0.4684 - acc: 0.8354 - val_loss: 0.6374 - val_acc: 0.7840

今度は6秒になりました。やはりImageDataGeneratorがネックとなっていそうです。
メモリが潤沢にある場合やサンプル数が沢山用意できる環境であればよいのですが、
大抵はそのような恵まれた環境ではないかと思います。
ImageDataGenerator自体を並列化することで対応してみたいと思います。

ImageDataGeneratorのプロセス並列化

実は、ImageDataGeneratorのflowメソッドがSequenceを継承したIteratorオブジェクトを返すので、プロセス並列化出来たりします。

ImageDataGeneratorのプロセス並列化をした場合

    model.fit_generator(datagen.flow(x_train, y_train,
                                     batch_size=batch_size),
                        epochs=epochs,
                        validation_data=(x_test, y_test),
                        workers=4)

    model.fit_generator(datagen.flow(x_train, y_train,
                                     batch_size=batch_size),
                        epochs=epochs,
                        validation_data=(x_test, y_test),
                        workers=32,
                        max_queue_size=64,
                        use_multiprocessing=True)

に変え、

data_augmentation = False

data_augmentation = True

に戻します。

結果

Epoch 90/100
196/196 [==============================] - 7s 37ms/step - loss: 0.7232 - acc: 0.7481 - val_loss: 0.6056 - val_acc: 0.7893
Epoch 91/100
196/196 [==============================] - 7s 36ms/step - loss: 0.7233 - acc: 0.7486 - val_loss: 0.5949 - val_acc: 0.7953
Epoch 92/100
196/196 [==============================] - 7s 36ms/step - loss: 0.7175 - acc: 0.7488 - val_loss: 0.5963 - val_acc: 0.7908
Epoch 93/100
196/196 [==============================] - 7s 36ms/step - loss: 0.7159 - acc: 0.7502 - val_loss: 0.6006 - val_acc: 0.7905
Epoch 94/100
196/196 [==============================] - 7s 36ms/step - loss: 0.7119 - acc: 0.7528 - val_loss: 0.6063 - val_acc: 0.7882
Epoch 95/100
196/196 [==============================] - 7s 36ms/step - loss: 0.7115 - acc: 0.7502 - val_loss: 0.5908 - val_acc: 0.7923
Epoch 96/100
196/196 [==============================] - 7s 36ms/step - loss: 0.7057 - acc: 0.7556 - val_loss: 0.5935 - val_acc: 0.7925
Epoch 97/100
196/196 [==============================] - 7s 36ms/step - loss: 0.6995 - acc: 0.7586 - val_loss: 0.5920 - val_acc: 0.7910
Epoch 98/100
196/196 [==============================] - 7s 36ms/step - loss: 0.6939 - acc: 0.7586 - val_loss: 0.5929 - val_acc: 0.7967
Epoch 99/100
196/196 [==============================] - 7s 37ms/step - loss: 0.6946 - acc: 0.7573 - val_loss: 0.5752 - val_acc: 0.8007
Epoch 100/100
196/196 [==============================] - 7s 37ms/step - loss: 0.6966 - acc: 0.7566 - val_loss: 0.5837 - val_acc: 0.7966

7秒まで早くなりました。並列化することでボトルネックをある程度解消できた事になります。

まとめ

複数GPUによる高速化は可能です。
ただし、fit_generatorを使う場合、generator部がボトルネックにならないように気をつける必要があります。

注意

今回提示したソースコードだと、modelのsave時にエラーになります。
原因はmulti_gpu_modelのDocに書いてあります。

# On model saving
    To save the multi-gpu model, use `.save(fname)` or `.save_weights(fname)`
    with the template model (the argument you passed to `multi_gpu_model`),
    rather than the model returned by `multi_gpu_model`.

したがって、モデルをsaveしたい時は、multi_gpu_model返すmodelではなく、
multi_gpu_model渡したmodelをsaveするようにしてください。

今回は本質と関係ない部分なので、対応は省略しました。