機械学習(Keras, PyTorch)と分散学習(Horovod)の環境構築シェルスクリプト

はじめに

• 機械学習や分散学習の環境構築は、AWSなら深層学習 AMIやSageMakerを使うのが便利です。
• 一方、GPUをバリバリ使う場合は、PC、GPU等の購入が良いと思います。
• 今回、Keras、PyTorchとHorovodの環境構築をAnsibleやTerraformで自動化しようと考えています。
• その前に、一連の処理をシェルスクリプト化して理解を深めました。
• 以下は、環境構築のシェルスクリプトの動作例です。だいたい、20分位で完了します。本文で詳細を説明します。

$ wget https://raw.githubusercontent.com/maedamikio/public/master/horovod/setup.sh
$ sh setup.sh
$ sudo reboot

Ubuntu Server 18.04.3 LTS

• クリーンインストールが前提です。

setup.sh

• シェルスクリプトの内容を解説します。

Version

• インストールするライブラリやパッケージ等のバージョンを指定しています。
• CUDAは、PyTorchやKerasのバックエンドのTensorFlowの推奨バージョン 10.0 にしています。ちなみに最新は 10.1 です。
• cuDNN、NCCL、TensorRTもCUDA 10.0に合わせています。
• Keras、TensorFlow、PyTorch、Open MPI、Horovodは、最新です。

CUDA=10.0.130-1
CUDNN=7.6.3.30-1+cuda10.0
NCCL=2.4.8-1+cuda10.0
NVINFER=5.1.5-1+cuda10.0
KERAS=2.2.5
TENSORFLOW=1.14.0
TORCH=1.2.0
TORCHVISION=0.4.0
OPENMPI=4.0.1
HOROVOD=0.18.1

Ubuntu

• Ubuntuのアップデート、アップグレードとシェルスクリプトに最低限必要なパッケージをインストールしています。

sudo apt update
sudo apt upgrade -y
sudo apt dist-upgrade -y
sudo apt install -y --no-install-recommends ca-certificates curl openssh-server wget
sudo apt autoremove -y

Docker

• Dockerのリポジトリの鍵を登録します。これで、aptでDockerのインストールができるようになります。
• 参考: https://docs.docker.com/install/linux/docker-ce/ubuntu/

sudo apt remove -y docker docker-engine docker.io containerd runc
sudo apt update
sudo apt install -y apt-transport-https ca-certificates curl gnupg-agent software-properties-common
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo apt-key add -
sudo add-apt-repository "deb [arch=amd64] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable"
sudo apt update
sudo apt install -y docker-ce docker-ce-cli containerd.io
sudo usermod -aG docker $USER

NVIDIA

• NVIDIAのリポジトリの鍵を登録します。これで、aptでCUDA等のインストールができるようになります。
• CUDA、cuDNN、NCCL、TensorRTのリポジトリのdebをダウンドードしてインストールします。
• ドライバ、CUDA、cuDNN、NCCL、TensorRTをバージョンを指定してインストールします。
• リポジトリやライブラリのバージョンが自動的にアップグレードしないようにします。
• 起動時にパーシステントモードにしています。これでドライバを常にロードしてくれるので、CUDA等の利用開始時の待ち時間がへります。
• 参考
  • https://www.tensorflow.org/install
  • https://hub.docker.com/r/nvidia/cuda/

sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-repo-ubuntu1804_10.0.130-1_amd64.deb
wget https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/nvidia-machine-learning-repo-ubuntu1804_1.0.0-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu1804_10.0.130-1_amd64.deb nvidia-machine-learning-repo-ubuntu1804_1.0.0-1_amd64.deb
sudo apt-mark hold cuda-repo-ubuntu1804 nvidia-machine-learning-repo-ubuntu1804
sudo apt update
sudo apt install -y --no-install-recommends nvidia-driver-430
sudo apt install -y --no-install-recommends cuda=$CUDA libcudnn7=$CUDNN libcudnn7-dev=$CUDNN libnccl-dev=$NCCL libnccl2=$NCCL libnvinfer-dev=$NVINFER libnvinfer5=$NVINFER
sudo apt-mark hold cuda libcudnn7 libcudnn7-dev libnccl-dev libnccl2 libnvinfer-dev libnvinfer5
echo "#!/bin/sh -e\nnvidia-smi -pm 1\nexit 0" | sudo tee /etc/init.d/nvidia && sudo chmod 755 /etc/init.d/nvidia && sudo ln -s ../init.d/nvidia /etc/rc5.d/S01nvidia

NVIDIA Docker

• NVIDIA Dockerのリポジトリの鍵を登録します。これで、aptでNVIDIA Dockerのインストールができるようになります。
• 参考: https://github.com/NVIDIA/nvidia-docker

distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt update
sudo apt install -y nvidia-container-toolkit

Python

• Python本体とpipを利用するため、最低限必要なパッケージをインストールしています。

sudo apt install -y --no-install-recommends python3 python3-dev python3-pip python3-setuptools

Machine learning

• 機械学習関連のパッケージをバージョンを指定してインストールしています。

sudo pip3 install keras==$KERAS tensorflow-gpu==$TENSORFLOW torch==$TORCH torchvision==$TORCHVISION

Open MPI

• ビルドに必要なパッケージのインストールしています。
• ソースファイルのダウンロード、解凍とフォルダの移動。
• Makefileの作成、ビルドとインストールをしています。
• ビルドは、$(nproc)でCPUコア数を指定して時間短縮をしています。

sudo apt install -y --no-install-recommends build-essential
wget https://www.open-mpi.org/software/ompi/v4.0/downloads/openmpi-$OPENMPI.tar.gz
tar zxf openmpi-$OPENMPI.tar.gz
cd openmpi-$OPENMPI
./configure
make -j $(nproc) all
sudo make install

Horovod

• Horovodのインストールは、TensorFlow、PyTorchやOpen MPIのインストールが完了している必要があります。
• HorovodでTensorFlowを使うには、TensorFlowのバイナリのg++が必要です。現在、g++-4.8となります。
• Remote Direct Memory Access関連のパッケージをインストールしています。これは、InfiniBand等の高速ネットワークを使う場合に利用します。
• Horovodインストール時に、NCCLの利用を指定しています。また、今回利用しないMXNetを除外しています。

sudo apt install -y --no-install-recommends build-essential g++-4.8 ibverbs-providers libibverbs1 librdmacm1
sudo HOROVOD_GPU_ALLREDUCE=NCCL HOROVOD_WITHOUT_MXNET=1 HOROVOD_WITH_PYTORCH=1 HOROVOD_WITH_TENSORFLOW=1 pip3 install horovod==$HOROVOD

SSH

• Horovodは、複数のマシンを利用する場合、Open MPIを利用します。
• その際、sshでの接続が発生します。その為、sshエージェントの転送やssh鍵の承認が必要になります。

sudo sed -i "s/#   ForwardAgent no/    ForwardAgent yes/" /etc/ssh/ssh_config
sudo sed -i "s/#   StrictHostKeyChecking ask/    StrictHostKeyChecking no/" /etc/ssh/ssh_config

MNIST

• 動作確認用のMNISTをダウンロードします。
• KerasとPyTorchの通常版。および、Horovod版です。

mkdir $HOME/examples
cd $HOME/examples
wget https://raw.githubusercontent.com/keras-team/keras/master/examples/mnist_cnn.py -O mnist_keras.py
wget https://raw.githubusercontent.com/pytorch/examples/master/mnist/main.py -O mnist_pytorch.py
wget https://raw.githubusercontent.com/horovod/horovod/master/examples/keras_mnist.py -O mnist_keras_horovod.py
wget https://raw.githubusercontent.com/horovod/horovod/master/examples/pytorch_mnist.py -O mnist_pytorch_horovod.py

リブート

• 最後にリブートします。

sudo reboot

Dockerfile

• Dockerfileの内容を解説します。
• 基本的には、コンテナに必要な箇所を上記のスクリプトを参考にしています。

FROM

• NVIDIA公式のイメージです。
• Ubuntu 18.04 に CUDA、cuDNNなどが含まれています。

FROM nvidia/cuda:10.0-cudnn7-devel-ubuntu18.04

Version

• シェルスクリプトと同様にバージョンの指定をしています。

ENV CUDA=10.0.130-1 \
    CUDNN=7.6.3.30-1+cuda10.0 \
    NCCL=2.4.8-1+cuda10.0 \
    NVINFER=5.1.5-1+cuda10.0 \
    KERAS=2.2.5 \
    TENSORFLOW=1.14.0 \
    TORCH=1.2.0 \
    TORCHVISION=0.4.0 \
    OPENMPI=4.0.1 \
    HOROVOD=0.18.1

Ubuntu, NVIDIA, Python, Open MPI, Horovod

• シェルスクリプトに記載のパッケージのうち、必要な物を一括でインストールしています。
• シェルスクリプトと同様にバージョンの指定をしています。

RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential \
    ca-certificates \
    curl \
    g++-4.8 \
    ibverbs-providers \
    libibverbs1 \
    librdmacm1 \
    libcudnn7=$CUDNN \
    libcudnn7-dev=$CUDNN \
    libnccl-dev=$NCCL \
    libnccl2=$NCCL \
    libnvinfer-dev=$NVINFER \
    libnvinfer5=$NVINFER \
    openssh-server \
    python3 \
    python3-dev \
    python3-pip \
    python3-setuptools \
    wget \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/* \
    && apt-mark hold \
    libcudnn7 \
    libcudnn7-dev \
    libnccl-dev \
    libnccl2 \
    libnvinfer-dev \
    libnvinfer5

Machine learning

• シェルスクリプトと同様にバージョンの指定をしています。

RUN pip3 install \
    keras==$KERAS \
    tensorflow-gpu==$TENSORFLOW \
    torch==$TORCH \
    torchvision==$TORCHVISION

Open MPI

• シェルスクリプトと同様です。ただし、データ削除の後処理をしています。

RUN mkdir /tmp/openmpi && \
    cd /tmp/openmpi && \
    wget https://www.open-mpi.org/software/ompi/v4.0/downloads/openmpi-$OPENMPI.tar.gz && \
    tar zxf openmpi-$OPENMPI.tar.gz && \
    cd openmpi-$OPENMPI && \
    ./configure && \
    make -j $(nproc) all && \
    make install && \
    rm -rf /tmp/openmpi

Horovod

• ldconfigでcuda関連のライブラリを指定しています。コンテナイメージのビルドには、必要な作業となります。

RUN ldconfig /usr/local/cuda/targets/x86_64-linux/lib/stubs && \
    HOROVOD_GPU_ALLREDUCE=NCCL \
    HOROVOD_WITHOUT_MXNET=1 \
    HOROVOD_WITH_PYTORCH=1 \
    HOROVOD_WITH_TENSORFLOW=1 \
    pip3 install horovod==$HOROVOD

SSH

• コンテナにsshでログインするための設定です。
• ユーザは、rootとしています。
• 公開鍵は、ホスト側の物を利用しています。

ARG AUTHORIZED_KEYS
RUN mkdir /var/run/sshd && \
    sed -i "s/#   ForwardAgent no/    ForwardAgent yes/" /etc/ssh/ssh_config && \
    sed -i "s/#   StrictHostKeyChecking ask/    StrictHostKeyChecking no/" /etc/ssh/ssh_config && \
    sed -i "s/#PermitRootLogin prohibit-password/PermitRootLogin yes/" /etc/ssh/sshd_config && \
    echo "root:pass" | chpasswd && \
    mkdir /root/.ssh && \
    echo "$AUTHORIZED_KEYS" > /root/.ssh/authorized_keys && \
    chmod 600 /root/.ssh/authorized_keys && \
    chown root:root /root/.ssh/authorized_keys

ENTRYPOINT、CMD

• コンテナ起動時にsshサーバを起動します。
• Horovodのコンテナのネットワークは、ホスト側を指定する必要があります。
• そのため、ホスト側のポート22とバッティングしないように10022を指定しています。
• ただし、コンテナ起動時に修正可能にしています。

ENTRYPOINT ["/usr/sbin/sshd", "-D", "-p"]
CMD ["10022"]

docker build

• temp等の適当なフォルダを作成して上記のDockerfileをダウンロードしています。
• ビルドのオプションでホームディレクトリ配下のssh公開鍵を指定しています。

mkdir temp
cd temp
wget https://raw.githubusercontent.com/maedamikio/public/master/horovod/Dockerfile
docker build -t horovod:latest --build-arg AUTHORIZED_KEYS="$(cat ~/.ssh/authorized_keys)" .

動作確認

• Keras、PyTorch、HorovodとDockerの各パターンでMNISTのテストをします。
• Horovodは、2ホストの例です。

Keras, PyTorch

$ cd ~/examples

$ python3 mnist_keras.py
Using TensorFlow backend.
省略
Epoch 12/12
60000/60000 [==============================] - 8s 129us/step - loss: 0.0265 - acc: 0.9916 - val_loss: 0.0270 - val_acc: 0.9914
Test loss: 0.027020889360253934
Test accuracy: 0.9914

$ cd ~/examples

$ python3 mnist_pytorch.py
Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to ../data/MNIST/raw/train-images-idx3-ubyte.gz
省略
Train Epoch: 10 [59520/60000 (99%)]	Loss: 0.007617

Test set: Average loss: 0.0318, Accuracy: 9892/10000 (99%)

Docker + Keras, PyTorch

ホストへのssh接続

• ホストへのssh接続時に -A を指定して、sshエージェントの転送を行ってください。わからない方は、検索等で調べてください。
• ホストへ接続後、sshエージェントが転送されているか確認します。

$ ssh -A ホスト名

$ ssh-add -l
4096 SH!256:xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx [e-mail] (RSA)

Dockerコンテナの起動

• -dでコンテナ起動後にデタッチしています。
• -vでホームフォルダをコンテナにマウントしています。今回は ubuntu ユーザーを使いました。つまり /home/ubuntu をマウントしています。
• --gpus allでGPUをコンテナにすべてマウントしています。個別に指定する事も可能です。詳しくは公式を参考にしてください。
• --nameでhorovodと名前を指定しています。お好みで適当に命名してください。
• --network=hostでコンテナのネットワークをホストと同居しています。Horovodを利用するために必要です。Horovodを利用しない場合は、不要です。今回は、下記でHorovodを利用するので手間を省くため指定しています。
• --rmでコンテナ終了時に削除するようにしています。お好みで省いてOKです。
• 最後にdocker buildで作成したhorovod:latestのイメージを指定しています。

$ docker run -d -v $HOME:$HOME --gpus all --name horovod --network=host --rm horovod

コンテナへのssh接続

• コンテナのsshポートを指定して接続します。
• rootでログインした事を確認します。

$ ssh root@localhost -p 10022

# pwd
/root

# cd /home/ubuntu/examples

# python3 mnist_keras.py
Using TensorFlow backend.
省略
Epoch 12/12
60000/60000 [==============================] - 8s 131us/step - loss: 0.0258 - acc: 0.9920 - val_loss: 0.0284 - val_acc: 0.9911
Test loss: 0.028411623225609948
Test accuracy: 0.9911

# python3 mnist_pytorch.py
Train Epoch: 1 [0/60000 (0%)]	Loss: 2.300039
省略
Train Epoch: 10 [59520/60000 (99%)]	Loss: 0.007857

Test set: Average loss: 0.0318, Accuracy: 9897/10000 (99%)

Horovod + Keras, PyTorch

• 2ホストの例です。
• 各ホストでsetup.shを実施している必要があります。
• 各ホストのIPアドレスは、192.168.1.2と192.168.1.3とした例となります。
• SSHエージェントが利用できる必要があります。

フォルダの確認

• 192.168.1.2の作業となります。
• -np でGPUの合計を指定します。今回は、各ホストに1GPU、合計2の場合を示しています。
• -H は 192.168.1.2 に 1GPU、192.168.1.3 に 1GPUがある事を意味します。
• pwdコマンドを実行してフォルダの位置が正しいか確認します。
• [1,0]や[1,1]でホスト毎の出力を識別できます。

$ cd examples

$ horovodrun -np 2 -H 192.168.1.2:1,192.168.1.3:1 pwd
[1,1]<stdout>:/home/ubuntu/examples
[1,0]<stdout>:/home/ubuntu/examples

Keras

• 2GPUを利用しているので、通常12エポックになっているのが、6エポックになっています。詳しくはコードの中身を確認してください。

$ horovodrun -np 2 -H 192.168.1.2:1,192.168.1.3:1 python3 mnist_keras_horovod.py
省略
[1,0]<stdout>:Epoch 6/6
[1,0]<stdout>:60000/60000 [==============================] - 14s 242us/step - loss: 0.0235 - acc: 0.9921 - val_loss: 0.0248 - val_acc: 0.9929
[1,1]<stdout>:Test loss: 0.024755005714310527
[1,1]<stdout>:Test accuracy: 0.9929
[1,0]<stdout>:Test loss: 0.024755004244462316
[1,0]<stdout>:Test accuracy: 0.9929

PyTorch

• こちらは、通常時と同様のエポックです。ただ、出力を確認すると両方のホストで実施されているのが確認できています。

$ horovodrun -np 2 -H 192.168.1.2:1,192.168.1.3:1 python3 mnist_pytorch_horovod.py
省略
[1,1]<stdout>:Train Epoch: 10 [29440/30000 (98%)]	Loss: 0.153234
[1,0]<stdout>:Train Epoch: 10 [29440/30000 (98%)]	Loss: 0.157580
[1,0]<stdout>:
[1,0]<stdout>:Test set: Average loss: 0.0564, Accuracy: 98.17%
[1,0]<stdout>:

Docker + Horovod + Keras, PyTorch

• 2ホストの例です。
• 各ホストでsetup.shを実施している必要があります。
• 各ホストのIPアドレスは、192.168.1.2と192.168.1.3とした例となります。
• SSHエージェントが利用できる必要があります。

Dockerコンテナの起動

• 各ホストでコンテナを起動してください。

$ docker run -d -v $HOME:$HOME --gpus all --name horovod --network=host --rm horovod

コンテナへのssh接続

• 192.168.1.2の作業となります。
• コンテナのsshポートを指定して接続します。
• rootでログインした事を確認します。

$ ssh root@192.168.1.2 -p 10022

# pwd
/root

# cd /home/ubuntu/examples

フォルダの確認

• 上記のホストでのHorovodとほぼ一緒です。
• ただし、コンテナのsshのポートを指定する必要があります。

$ cd examples

$ horovodrun -np 2 -H 192.168.1.2:1,192.168.1.3:1 -p 10022 pwd
[1,1]<stdout>:/home/ubuntu/examples
[1,0]<stdout>:/home/ubuntu/examples

Keras

• コンテナのsshポートを指定している部分が異なります。

$ horovodrun -np 2 -H 192.168.1.2:1,192.168.1.3:1 -p 10022 python3 mnist_keras_horovod.py
省略
[1,0]<stdout>:Epoch 6/6
[1,0]<stdout>:60000/60000 [==============================] - 14s 242us/step - loss: 0.0236 - acc: 0.9923 - val_loss: 0.0287 - val_acc: 0.9916
[1,1]<stdout>:Test loss: 0.028658814687919677
[1,1]<stdout>:Test accuracy: 0.9916
[1,0]<stdout>:Test loss: 0.028658814687919677
[1,0]<stdout>:Test accuracy: 0.9916

PyTorch

• コンテナのsshポートを指定している部分が異なります。

$ horovodrun -np 2 -H 192.168.1.2:1,192.168.1.3:1 -p 10022 python3 mnist_pytorch_horovod.py
省略
1,0]<stdout>:Train Epoch: 10 [29440/30000 (98%)]	Loss: 0.151875
[1,1]<stdout>:Train Epoch: 10 [29440/30000 (98%)]	Loss: 0.163286
[1,0]<stdout>:
[1,0]<stdout>:Test set: Average loss: 0.0571, Accuracy: 98.15%
[1,0]<stdout>:

KerasとKeras Horovodの差分の解説

• 通常のKerasをほんの少し修正することで、Horovodに対応できます。

• PyTorchに関しては、別途解説できればと思います。

• 本体で利用するライブラリをインポートしています。

> import math
> import tensorflow as tf
> import horovod.keras as hvd

• Horovodの初期化とconfig設定をしています。

> 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))

• エポックをGPUの数に合わせて修正しています。

< epochs = 12
> epochs = int(math.ceil(12.0 / hvd.size()))

• オプティマイザーをGPUの数に合わせて修正し、ラッピングしています。

< 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'))

• 学習時のコールバックに上記のチェックポイントを指定しています。

< 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 if hvd.rank() == 0 else 0, validation_data=(x_test, y_test))

おわりに

• 環境構築の簡便化のためのシェルスクリプトとDockerfileの説明をしました。
• Horovodの利用は、複数ホストになるため、NFS等の共有ストレージが必要となるでしょう。
• また、複数ホストでの認証系も統一する必要があるかもしれません。ただ、個人で利用するのであれば、ubuntu等で統一すればOKです。
• 今後は、KubernetesのGPUスケジューリング、AnsibleやTerraformなどを考慮しながら検討を進めたいと思います。