- はじめに -

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

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

- 前提 -

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

github.com

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

また、公式によると普通に書くより早いらしい（未検証）

複数ノードで利用する場合、各ノードが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 osimport tracebackimport datetimeimport torchimport torchvision.transformsas transformsimport torch.nnas nnimport torch.optimas optimfrom torch.utils.dataimport Datasetfrom cnn_finetuneimport make_modelimport pandasas pdfrom PILimport Image# --- 追加 ---import horovod.torchas hvdtorch.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()classMyDataSet(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):returnlen(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())# ----------deftrain(epoch):    total_loss =0    total_size =0    model.train()    train_sampler.set_epoch(epoch)for batch_idx, (data, target)inenumerate(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 ==0and hvd.rank() ==0:            now = datetime.datetime.now()withopen('/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))# maintry:for epochinrange(1,100):        train(epoch)# --- hvd.rankでstate_dict保存を絞るよう変更 ---if hvd.rank() ==0:            torch.save(model.state_dict(),'/mnt/senet154_{}.model'.format(epoch))exceptExceptionas e:    now = datetime.datetime.now()withopen('/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\-xPATH=$PATH-xPYTHONPATH=$PYTHONOATH-xLD_LIBRARY_PATH=$LD_LIBRARY_PATH-xCPATH=$CPATH-xLIBRARY_PATH=$LIBRARY_PATH-xNCCL_ROOT=$NCCL_ROOT\-bind-to none\-map-by slot\--hostfile /mnt/host.txt\-np8\     python3 /mnt/train.py

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

• -mca pml ob1
• -mcabtl ^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の方法を追加できたりするので結構柔軟に書けるとも思う。

要は使い分け。