HorovodとKerasを利用したMNIST CNNの分散学習

はじめに

• 元々、分散処理に興味があります。キャリアの始めはGnutella等のP2Pでした。
• と言う訳で、機械学習の分散処理の「分散学習」にも注目しています。
• 現在、Uber社のHorovodが注目されているので試してみました。
• あと、Kubeflowと連携出来るのは、ポイント高いですね。

EC2インスタンスの作成

• AMI: Deep Learning Base AMI (Ubuntu) Version 11.0 - ami-081943b954855ddfe
• インスタンス: p2.8xlarge
• CPU: Intel(R) Xeon(R) CPU E5-2686 v4 @ 2.30GHz x 8コア
• MEM: 488G
• GPU: Tesla K80 x 8基
• NW: 10Gbps
• SSH鍵: id_rsa
• 台数: 2台

IPアドレス

• アドレスは、状況に応じて適当に読み換える

インスタンス1

• パブリック: 18.182.44.122
• プライベート: 10.255.254.195

インスタンス2

• パブリック: 18.182.44.60
• プライベート: 10.255.254.76

SSH接続等

• Horovodは、Open MPIを使うので、事前にSSHの接続確認をする

ssh鍵の登録

$ ssh-add .ssh/id_rsa

インスタンス1

# パブリックから接続
$ ssh ubuntu@18.182.44.122

# インスタンス2のプライベートへ接続
$ ssh ubuntu@10.255.254.76

# インスタンス1に戻る
$ exit

インスタンス2

# パブリックから接続
$ ssh ubuntu@18.182.44.60

# インスタンス1のプライベートへ接続
$ ssh ubuntu@10.255.254.195

# インスタンス2に戻る
$ exit

Horovodのインストール

• horovodの依存関係で、keras、tensorflow-gpuとtorchを最初にインストールしている
• インスタンス1とインスタンス2で実施

$ pip3 install --user keras tensorflow-gpu torch
$ pip3 install --user horovod

mpirunのためのhostfileの作成

• slotsはGPUの数
• 比較のため複数のパターンを作成
• インスタンス1で実施

$ vim hostfile-1GPU
10.255.254.195 slots=1

$ vim hostfile-8GPU
10.255.254.195 slots=8

$ vim hostfile-16GPU
10.255.254.195 slots=8
10.255.254.76  slots=8

mpirunの動作確認

• hostnameの分散処理を実施
• インスタンス1で実施
• npはGPUの数を表す

$ mpirun -np 16 --hostfile hostfile-16GPU hostname
ip-10-255-254-195
ip-10-255-254-195
ip-10-255-254-195
ip-10-255-254-195
ip-10-255-254-195
ip-10-255-254-195
ip-10-255-254-195
ip-10-255-254-195
ip-10-255-254-76
ip-10-255-254-76
ip-10-255-254-76
ip-10-255-254-76
ip-10-255-254-76
ip-10-255-254-76
ip-10-255-254-76
ip-10-255-254-76

KerasとHorovodのMNIST CNNをGitHubのexamplesからダウンロード

• インスタンス1とインスタンス2で実施

# Keras
$ wget https://raw.githubusercontent.com/keras-team/keras/master/examples/mnist_cnn.py

# Horovod
$ wget https://raw.githubusercontent.com/uber/horovod/master/examples/keras_mnist.py

mnist_cnn.pyとkeras_mnist.pyの差分

• コメント等や改行等を修正済み
• Horovodの初期化やGPU数に応じた各種修正が行われている

$ diff mnist_cnn.py keras_mnist.py
7a8,17
> import math
> import tensorflow as tf
> import horovod.keras as hvd
>
> hvd.init()
>
> config = tf.ConfigProto()
> config.gpu_options.allow_growth = True
> config.gpu_options.visible_device_list = str(hvd.local_rank())
> K.set_session(tf.Session(config=config))
11c21,22
< epochs = 12
---
>
> epochs = int(math.ceil(12.0 / hvd.size()))
49c60,69
< model.compile(loss=keras.losses.categorical_crossentropy, optimizer=keras.optimizers.Adadelta(), metrics=['accuracy'])
---
> opt = keras.optimizers.Adadelta(1.0 * hvd.size())
>
> opt = hvd.DistributedOptimizer(opt)
>
> model.compile(loss=keras.losses.categorical_crossentropy, optimizer=opt, metrics=['accuracy'])
>
> callbacks = [hvd.callbacks.BroadcastGlobalVariablesCallback(0),]
>
> if hvd.rank() == 0:
>     callbacks.append(keras.callbacks.ModelCheckpoint('./checkpoint-{epoch}.h5'))
51c71
< model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, verbose=1, validation_data=(x_test, y_test))
---
> model.fit(x_train, y_train, batch_size=batch_size, callbacks=callbacks, epochs=epochs, verbose=1, validation_data=(x_test, y_test))

MNIST画像のダウンロードのための試験走行

• インスタンス1とインスタンス2で実施

$ /usr/bin/python3 mnist_cnn.py
省略
Test loss: 0.024124369315993682
Test accuracy: 0.9917

batch_sizeとepochsの修正

• 16GPUだとあっと言う間に処理が終わってしまうため、処理時間の比較のため以下の様に修正

mnist_cnn.py

# batch_size = 128
batch_size = 10000

# epochs = 12
epochs = 200

keras_mnist.py

# batch_size = 128
batch_size = 10000

# epochs = int(math.ceil(12.0 / hvd.size()))
epochs = int(math.ceil(200.0 / hvd.size()))

別端末からGPU使用率を確認する

インスタンス1

$ ssh ubuntu@18.182.44.122

$ watch -n 0.1 nvidia-smi

インスタンス2

$ ssh ubuntu@18.182.44.60

$ watch -n 0.1 nvidia-smi

MNIST CNNの実行

• インスタンス1で実行

Keras 1GPU

$ time /usr/bin/python3 mnist_cnn.py
省略
Test loss: 0.02959173348431377
Test accuracy: 0.9925

real	17m10.213s
user	7m59.740s
sys	3m22.208s

Horovod 1GPU

$ time mpirun -np 1 \
    --hostfile hostfile-1GPU \
    -bind-to none -map-by slot \
    -x NCCL_DEBUG=INFO -x LD_LIBRARY_PATH -x PATH \
    -x NCCL_SOCKET_IFNAME=^docker0 \
    -mca pml ob1 -mca btl ^openib \
    -mca btl_tcp_if_exclude lo,docker0 \
    /usr/bin/python3 keras_mnist.py
省略
Test loss: 0.029745910675377855
Test accuracy: 0.9909

real	19m18.805s
user	8m29.088s
sys	3m48.888s

Horovod 8GPU

$ time mpirun -np 8 \
    --hostfile hostfile-8GPU \
    -bind-to none -map-by slot \
    -x NCCL_DEBUG=INFO -x LD_LIBRARY_PATH -x PATH \
    -x NCCL_SOCKET_IFNAME=^docker0 \
    -mca pml ob1 -mca btl ^openib \
    -mca btl_tcp_if_exclude lo,docker0 \
    /usr/bin/python3 keras_mnist.py
省略
Test loss: 0.03558127259654575
Test accuracy: 0.9872

real	2m45.823s
user	8m20.308s
sys	0m45.080s

Horovod 16GPU

$ time mpirun -np 16 \
    --hostfile hostfile-16GPU \
    -bind-to none -map-by slot \
    -x NCCL_DEBUG=INFO -x LD_LIBRARY_PATH -x PATH \
    -x NCCL_SOCKET_IFNAME=^docker0 \
    -mca pml ob1 -mca btl ^openib \
    -mca btl_tcp_if_exclude lo,docker0 \
    /usr/bin/python3 keras_mnist.py
省略
Test loss: 0.10731917524505406
Test accuracy: 0.9643

real	1m36.460s
user	3m7.332s
sys	4m25.624s

結果一覧

項目	Test loss	Test accuracy	real	user	sys
Keras 1GPU	0.02959173348431377	0.9925	17m10.213s	7m59.740s	3m22.208s
Horovod 1GPU	0.029745910675377855	0.9909	19m18.805s	8m29.088s	3m48.888s
Horovod 8GPU	0.03558127259654575	0.9872	2m45.823s	8m20.308s	0m45.080s
Horovod 16GPU	0.10731917524505406	0.9643	1m36.460s	3m7.332s	4m25.624s

おわりに

• EC2のDeep Learning Base AMIを利用する事で、マルチGPUとマルチノードの分散学習を簡単に実施出来ることを確認しました。
• 本当は、TensorFlow、KerasとHorovodが準備済みのDeep Learning AMIを使いたかったのですが、sshログイン時のAnacondaの自動アクティベートが分からなかったため、上記の手順になりました。
• 時間の短縮は、期待通りと思っています。ただし、精度は若干悪くなっていますね。この辺りは、専門でないので、割愛します。