Stimulator

機械学習とか好きな技術話とかエンジニア的な話とかを書く

horovodを用いたPytorchの分散学習

Python スクリプト 機械学習

- はじめに -

近年、分散深層学習の研究、ライブラリ開発が盛んに行われている。

本記事はuber社が公開しているhorovodを利用した分散CNNのメモである。


　

- 前提 -

horovodとは、バックエンドをOpenMPIとしTensorFlow、Keras、PyTorchを最小限のコード変更で分散学習できるようにするためのパッケージである。

github.com

現状TensorFlowを使って書かれたコードをDistributed TensorFlowに対応させるにはパラメータサーバやマスタサーバの動きを理解した上で多くの変更を要するが、horovodではそれらをncclのall reduceを利用しwrappingしてあるため、最小限のコード変更で分散学習が可能となる。

また、公式によると普通に書くより早いらしい（未検証）
https://user-images.githubusercontent.com/16640218/38965607-bf5c46ca-4332-11e8-895a-b9c137e86013.png


複数ノードで利用する場合、各ノードがOpenMPIを通して疎通できる必要がある。その環境構築については以下に記載している。ChainerMNが動けば、ほぼ変更なくhorovodを動かす事ができる。

vaaaaaanquish.hatenablog.com

OpenMPI周りの設定が終わったらpipでhorovodを導入する。

pip install horovod

もしくは、DockerHubにHorovod-dockerも公開されていため、バックエンドの設定が整えば、こちらを利用する事で分散学習を始められる。
horovod/docker.md at master · uber/horovod · GitHub


　
PyTorchでCNNモデルを簡易に利用する方法は以下に記載している。
以下に記載のpretrain modelを利用したCNNモデルをhorovodで分散学習させる。

vaaaaaanquish.hatenablog.com


　

- 学習スクリプトの変更 -

学習を行うtrain.pyを以下に示す。
学習スクリプトはpretrainを学習させる記事に詳細を書いてあるので参考に。
読み込むデータのPathやログ出力先はマウントしているディレクトリ等でよしなに。

（※ 以下はPyTorch 0.4.0ですが、バージョンによってDataloader周りとかちょいちょい違いがあるので注意）

import os
import traceback
import datetime
import torch
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset
from cnn_finetune import make_model
import pandas as pd
from PIL import Image

# --- 追加 ---
import horovod.torch as hvd
torch.manual_seed(42)
hvd.init()
torch.cuda.set_device(hvd.local_rank())
torch.cuda.manual_seed(42)
# ----------

# 10クラス分類を想定
model = make_model('senet154', num_classes=10, pretrained=True, input_size=(256, 256))
criterion = nn.CrossEntropyLoss()

class MyDataSet(Dataset):
    def __init__(self, csv_path, root_dir):
        self.train_df = pd.read_csv(csv_path)
        self.root_dir = root_dir
        self.images = os.listdir(self.root_dir)
        # normalizeのmean, stdはpretrain modelより
        # https://github.com/Cadene/pretrained-models.pytorch/tree/master/pretrainedmodels/models
        self.transform = transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ColorJitter(brightness=1, contrast=1, saturation=1, hue=0.5),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        
    def __len__(self):
        return len(self.images)
    
    def __getitem__(self, idx):
        image_name = self.images[idx]
        image = Image.open( os.path.join(self.root_dir, image_name) )
        image = image.convert('RGB')
        label = self.train_df.query('ImageName=="'+image_name+'"')['ImageLabel'].iloc[0]
        return self.transform(image), int(label)

train_set = MyDataSet('train.csv', './train')

# --- 追加, 変更 ---
train_sampler = torch.utils.data.distributed.DistributedSampler(
    train_set, num_replicas=hvd.size(), rank=hvd.rank())
train_loader = torch.utils.data.DataLoader(
    train_set, batch_size=batch_size, sampler=train_sampler, pin_memory=True)
hvd.broadcast_parameters(model.state_dict(), root_rank=0)
model.cuda()
optimizer = optim.SGD(model.parameters(), lr=0.01 * hvd.size(),
                      momentum=0.9)
optimizer = hvd.DistributedOptimizer(
    optimizer, named_parameters=model.named_parameters())
# ----------

    
def train(epoch):
    total_loss = 0
    total_size = 0
    model.train()
    train_sampler.set_epoch(epoch)

    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.cuda(), target.cuda()
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        total_loss += loss.item()
        total_size += data.size(0)
        loss.backward()
        optimizer.step()

        # --- hvd.rankで出力を絞るよう変更 ----
        if batch_idx % 100 == 0 and hvd.rank() == 0:
            now = datetime.datetime.now()
            with open('/mnt/log.text', 'a') as fa:
                fa.write('[{}] Train Epoch: {} [{}/{} ({:.0f}%)]\tAverage loss: {:.20f}\n'.format(now, epoch, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader), total_loss / total_size))


# main
try:
    for epoch in range(1, 100):
        train(epoch)
         # --- hvd.rankでstate_dict保存を絞るよう変更 ---
        if hvd.rank() == 0:
            torch.save(model.state_dict(), '/mnt/senet154_{}.model'.format(epoch))
except Exception as e:
    now = datetime.datetime.now()
    with open('/mnt/log.text', 'a') as fa:
        fa.write('[{}]error: {}\n'.format(now, str(e)))
        fa.write(traceback.format_exc()+'\n')
    raise

変更箇所はコメントの通り少ない。

初回起動時に「cnn_finetune.make_model」でpretrainモデルのダウンロードが走ってしまうため、複数ノードでdockerを利用するなら一回docker内でmake_modelを実行しダウンロードしてモデルファイルをdocker imageに含めるか、lusterfsのような共通で見れるディレクトリをマウントしてそのモデルを参照するようにすると良い。

また、こちらで学習したモデル(state_dict)は、前述したpretrainを学習させる記事内のtestコードで推論できる。


　

- 実行 -

動作させるどこかしらのノードないしdockerにログインし以下を実行する。

もし各ノードでdockerを利用している場合は、PyTorchではdocker run時に「--ipc=host」を付けなければ「Unable to write to file」となってしまう事に留意する。
Unable to write to file </torch_18692_1954506624> - PyTorch Forums

また、horovodでは/tmpを利用するため、docker run時に「-v /tmp:/tmp」等としtmpもマウントしておく必要がある。

　
mpiexecコマンドを利用し実行する。
hostfileについてはChainerMNの記事参照。

mpiexec --allow-run-as-root \
     --mca btl_tcp_if_include ib0 \
     -mca pml ob1 \
     -mca btl ^openib \
     -x PATH=$PATH -x PYTHONPATH=$PYTHONOATH -x LD_LIBRARY_PATH=$LD_LIBRARY_PATH -x CPATH=$CPATH -x LIBRARY_PATH=$LIBRARY_PATH -x NCCL_ROOT=$NCCL_ROOT \
     -bind-to none \
     -map-by slot \
     --hostfile /mnt/host.txt \
     -np 8 \
     python3 /mnt/train.py

dockerを利用している場合「--allow-run-as-root 」が必須である。
また、ChainerMNの記事に記載のコマンドとの違いとして、以下を設定しTCP通信を強制する必要がある。

  • -mca pml ob1
  • -mca btl ^openib

これを設定しないと以下のようにsubprocessが次々死んでいき、全体の動作も止まってしまう。

HorovodBroadcast_residual_layer_batch_normalization_moving_variance_0 [missing ranks: 1]

　
OpenMPI 3以降であれば、以下を利用してprocessを単一CPUにバインドさせないようにする。
また、defaultではNUMA設定が単一となってしまうため、map-by slotも利用しておくと良いらしい。

  • -bind-to none
  • -map-by slot


ログが吐かれ始めれば成功。


参考：https://github.com/uber/horovod/blob/master/docs/running.md


　

- おわりに -

「horovodならコード変更最小限に分散学習！」とは言うけど、OpenMPIが動く前提があり、正直「何よりOpenMPIが動作する環境を作るのがしんどいんじゃい…」と思う。

OpenMPIのsettingが一通り上手くいってしまえば、後はかなり自由にモデリングできると思う。各ノードにhvd.broadcastで別々のデータを送ったり、hvd.allreduceでなくallgatherを使えばaggregationの方法を追加できたりするので結構柔軟に書けるとも思う。

要は使い分け。



　

vaaaaaanquish

関連記事