- 前提 -
horovodとは、バックエンドをOpenMPIとしTensorFlow、Keras、PyTorchを最小限のコード変更で分散学習できるようにするためのパッケージである。
現状TensorFlowを使って書かれたコードをDistributed TensorFlowに対応させるにはパラメータサーバやマスタサーバの動きを理解した上で多くの変更を要するが、horovodではそれらをncclのall reduceを利用しwrappingしてあるため、最小限のコード変更で分散学習が可能となる。
また、公式によると普通に書くより早いらしい（未検証）
複数ノードで利用する場合、各ノードがOpenMPIを通して疎通できる必要がある。その環境構築については以下に記載している。ChainerMNが動けば、ほぼ変更なくhorovodを動かす事ができる。
OpenMPI周りの設定が終わったらpipでhorovodを導入する。
pip install horovod
もしくは、DockerHubにHorovod-dockerも公開されていため、バックエンドの設定が整えば、こちらを利用する事で分散学習を始められる。
horovod/docker.md at master · uber/horovod · GitHub
PyTorchでCNNモデルを簡易に利用する方法は以下に記載している。
以下に記載のpretrain modelを利用したCNNモデルをhorovodで分散学習させる。
- 学習スクリプトの変更 -
学習を行う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の方法を追加できたりするので結構柔軟に書けるとも思う。
要は使い分け。
