HorovodでTensorflowの分散トレーニングについて

まいど、SBCloud エンジニアのKou(寇)です。

最新のディープラーニング モデルのトレーニングには、多くの場合、GPU によって提供される大量の計算が必要です。1 つの GPU から多くの GPU へのスケーリングにより、トレーニングと研究のプロセスを大幅に短縮できます。この記事ではAlibaba CloudのマルチGPUのインスタンスを利用して、Horovodなどの簡単なセットアップとデモを紹介します。

Horovodとは

Horovodは、2017年Uberが開発したTensorFlowのオープンソース分散トレーニングフレームワークです。その目標は、分散ディープラーニングを集約処理(オールリデュース)を介して迅速かつ使いやすくするものであり、通常のtensorflow実装に対して、少ない変更量で分散実行できるようになるのは優れたメリットです。

github.com

環境構築

1. Alibaba CloudのGPUインスタンスを購入する

  • リージョン:東京
  • イメージ:Ubuntu 16.04 64bit
  • ストレージ:Ultraクラウドディスク40GB
  • インスタンスタイプ:ecs.gn5-c4g1.2xlarge(8vcpu、60GBRAM)
  • CPUスペック:Intel Xeon E5-2682v4(2.5GHz)
  • GPUスペック:NVIDIA P100*2 (3584 Pascal CUDA Core、12GB GDDR5 ビデオメモリ、9.3TFlops single、4.7T​​Flops double)

f:id:sbc_kou:20190910101901p:plain

2. CUDAの設定

現時点(2019/9/10)の最新バージョン(CUDA Toolkit 10.1 )のリンクからToolkitをダウンロードしてインストールします。

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1604/x86_64/cuda-ubuntu1604.pin
sudo mv cuda-ubuntu1604.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget http://developer.download.nvidia.com/compute/cuda/10.1/Prod/local_installers/cuda-repo-ubuntu1604-10-1-local-10.1.243-418.87.00_1.0-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu1604-10-1-local-10.1.243-418.87.00_1.0-1_amd64.deb
sudo apt-key add /var/cuda-repo-10-1-local-10.1.243-418.87.00/7fa2af80.pub
sudo apt-get update
sudo apt-get -y install cuda

3. cuDNNの設定

cuDNNは、ディープニューラルネットワーク用のGPUアクセラレーションライブラリであり、フォワードおよびバックレコンボリューション、プーリング、正規化などの標準ルーチンに対しての実装を提供します。本記事で最新のバージョンv7.6.3を使いました。

wget https://developer.nvidia.com/compute/machine-learning/cudnn/secure/7.6.3.30/Production/10.1_20190822/cudnn-10.1-linux-x64-v7.6.3.30.tgz
tar -xzvf cudnn-9.0-linux-x64-v7.tgz
sudo cp cuda/include/cudnn.h /usr/local/cuda/include
sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64
sudo chmod a+r /usr/local/cuda/include/cudnn.h /usr/local/cuda/lib64/libcudnn*

4. NCCLの設定

NCCLはNVIDIAが提供しているマルチGPU向けの集合通信用のライブラリです、下記のコマンドで最新のバージョンをダウンロードします。

wget https://developer.nvidia.com/compute/machine-learning/nccl/secure/v2.4/prod//nccl_2.4.8-1%2Bcuda10.1_x86_64.txz
cd /usr/local 
tar xvf nccl_2.4.8-1+cuda10.1_x86_64.txz

6. tensorflowの設定

ここではAnaconda仮想環境での設定方法を使いました。PyPIより、Anacondaを利用することで、tensorflow自体から、それと関連するライブラリが一括インストールされるので、とても便利です。

conda create -n tf pip python=3.7
source activate tf
conda install -c anaconda tensorflow-gpu 

7. MPIの設定

MPIは、並列コンピューティング利用するための標準化された規格であり、ここでMPIのオープンソース版Open MPIを使いました。

wget https://download.open-mpi.org/release/open-mpi/v3.1/openmpi-3.1.2.tar.bz2
gunzip -c openmpi-3.1.2.tar.gz | tar xf -
cd openmpi-4.0.1
 ./configure --prefix=/usr/local
make all install

8. Horovodのインストール

ようやく最後となりますが、下記のコマンドでhorovodをインストールします。

HOROVOD_GPU_ALLREDUCE=NCCL pip install --no-cache-dir horovod

Horovodの実装

本記事は、horovod検証の目的から、tensorflowの公式サイトに紹介されたシンプルな4層のCNNモデルを用いて、MNISTデータセット内の手書きの数字を認識する例を使いました。またマルチGPUに並列処理させる為に、データセットを分割して複数のGPU間でデータセットを準備する必要があります。

下記ソースコードのコメントアウトのように、hovorodモジュールをインポートして、ただ数箇所を変更して簡単に済むことができます。このような手間を削減するのは、特に開発エンジニアやデータサイエンティストにとっては楽となります。

import os
import tensorflow as tf
import horovod.tensorflow as hvd
import numpy as np
import time
from tensorflow import keras

layers = tf.layers
tf.logging.set_verbosity(tf.logging.INFO)

def conv_model(feature, target, mode):

    target = tf.one_hot(tf.cast(target, tf.int32), 10, 1, 0)
    feature = tf.reshape(feature, [-1, 28, 28, 1])

    with tf.variable_scope('conv_layer1'):
        h_conv1 = layers.conv2d(feature, 32, kernel_size=[5, 5],
                                activation=tf.nn.relu, padding="SAME")
        h_pool1 = tf.nn.max_pool(
            h_conv1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')

    with tf.variable_scope('conv_layer2'):
        h_conv2 = layers.conv2d(h_pool1, 64, kernel_size=[5, 5],
                                activation=tf.nn.relu, padding="SAME")
        h_pool2 = tf.nn.max_pool(
            h_conv2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
        h_pool2_flat = tf.reshape(h_pool2, [-1, 7 * 7 * 64])

    h_fc1 = layers.dropout(
        layers.dense(h_pool2_flat, 1024, activation=tf.nn.relu),
        rate=0.5, training=mode == tf.estimator.ModeKeys.TRAIN)

    logits = layers.dense(h_fc1, 10, activation=None)
    loss = tf.losses.softmax_cross_entropy(target, logits)
    tf.summary.scalar('loss', loss)

    correct_prediction = tf.equal(tf.argmax(logits, 1),
                                  tf.argmax(target, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    tf.summary.scalar('accuracy', accuracy)

    return tf.argmax(logits, 1), loss, accuracy

def train_input_generator(x_train, y_train, batch_size=64):
    assert len(x_train) == len(y_train)
    while True:
        p = np.random.permutation(len(x_train))
        x_train, y_train = x_train[p], y_train[p]
        index = 0
        while index <= len(x_train) - batch_size:
            yield x_train[index:index + batch_size], \
                  y_train[index:index + batch_size],
            index += batch_size

def main(_):
    start_time = time.time()
    
    '''---------------------- horovodの初期化 ----------------------'''
    hvd.init()

    (x_train, y_train), (x_test, y_test) = \
        keras.datasets.mnist.load_data('MNIST-data-%d' % hvd.rank())


    x_train = np.reshape(x_train, (-1, 784)) / 255.0
    x_test = np.reshape(x_test, (-1, 784)) / 255.0

    with tf.name_scope('input'):
        image = tf.placeholder(tf.float32, [None, 784], name='image')
        label = tf.placeholder(tf.float32, [None], name='label')
    predict, loss, accuracy = conv_model(image, label, tf.estimator.ModeKeys.TRAIN)

    '''------------- GPUの数に基づいて学習速度を調整する --------------'''
    opt = tf.train.AdamOptimizer(0.001 * hvd.size())

    '''------------- horovod分散オプティマイザを追加する--------------'''
    opt = hvd.DistributedOptimizer(opt)

    global_step = tf.train.get_or_create_global_step()
    train_op = opt.minimize(loss, global_step=global_step)

    hooks = [
        '''------------- 各GPU上のモデルを同じ値で初期化 --------------'''
        hvd.BroadcastGlobalVariablesHook(0),
        tf.train.StopAtStepHook(last_step=20000 // hvd.size()),
        tf.train.LoggingTensorHook(tensors={'step': global_step,
                       'loss': loss,
                       'accuracy':accuracy},
                                   every_n_iter=10),
    ]

    '''------------- 各プロセスにGPUを割り当てる --------------'''
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    config.gpu_options.visible_device_list = str(hvd.local_rank())

 '''-------- チェックポイントをワーカー0にのみ保存する---------'''
    checkpoint_dir = './new/checkpoints' if hvd.rank() == 0 else None
    training_batch_generator = train_input_generator(x_train,
                                                     y_train, batch_size=100)

    merged = tf.summary.merge_all()
    with tf.train.MonitoredTrainingSession(checkpoint_dir=checkpoint_dir,
                                           hooks=hooks,
                                           config=config) as mon_sess:
         writer = tf.summary.FileWriter("./tflog/%d" % hvd.rank(), mon_sess.graph)
         while not mon_sess.should_stop():
             image_, label_ = next(training_batch_generator)
             _, result, step = mon_sess.run([train_op,merged,global_step], feed_dict={image: image_, label: label_})
             if step % 100 == 0:
                 writer.add_summary(result,step)
    duration = time.time()-start_time
    print("/device:GPU:" + str(hvd.local_rank()) + "runtime is --- %s seconds ---" % duration)

if __name__ == "__main__":
    tf.app.run()

実行

2つのGPUで同時に実行させるため、下記の命令でtrain.pyスクリプトを実行します。 f:id:sbc_kou:20190930104204p:plain

mpirun -np 2 -bind-to none -map-by slot -x NCCL_DEBUG=INFO -x LD_LIBRARY_PATH -x PATH -mca pml ob1 -mca btl ^openib --allow-run-as-root python train.py

上記のスクリプト実行の途中に、別セッションでnvidia-smiコマンドでGPUの利用状況を確認すると、下図の赤枠に示したように、gpu device 0とgpu device 1、それぞれGPUの使用率が45%であることを分かりました。

f:id:sbc_kou:20190910152512p:plain

実行完了したら、下図の赤枠の部分から、gpu 0とgpu 1の実行時間を確認することができます。

f:id:sbc_kou:20190910152837p:plain

シングルGPUの結果も確認しました。下図のように、実行時間がマルチGPUよりそれほど差がないですが、gpuの数が増えるほど、horovodの効果が出ると思います。 f:id:sbc_kou:20190930104640p:plain

Tensorboardでログデータを確認しましょう。下記の命令を実行して、ブラウザでサーバの6006ポートに接続します。(事前にセキュリティグループのポートを開放する必要があります。)

tensorboard --logdir /logs/

gpu 0とgpu 1、10KのStepで99%以上の認識率を確認できました。またstepの回数が多くなるにつれて、Lossが減るのが見えます。

f:id:sbc_kou:20190910133912p:plain

最後

いかがでしたでしょうか

本記事ではAlibaba Cloud マルチGPUを持つインスタンスを利用してHorovodの実現性を検証しました。Horovodを用いてクラスタでの実行とかベンチマークの検証などの内容については、今後の記事で共有したいと思います。 また、従量課金制でAlibaba CloudのGPUインスタンスをお気軽にお試しできるので、ぜひためしてみてください!