ホリケン's diary

趣味はでぃーぷらーにんぐ

PyTorchにおける並列化処理について~備忘録~

動機
  • cpuの並列処理+GPUの並列処理が必要なモデルを実装する上で知識の整理をしておきたかったから

  • 時短という意味でもこれから使いたいから

  • 知らないことが多そうで単純に面白そうだったから

CPUでの処理並列化

(参照: Multiprocessing best practices — PyTorch master documentation)

import torch.multiprocessing as mp

num_processes = 4
model = MyModel()
model.share_memory()
processes = []
for rank in range(num_processes):
    p = mp.Process(target=train, args=(model,))
    p.start()
    processes.append(p)
for p in processes:
   p.join()

1~2行目でプロセス数の指定model()の読み込みを行なっている。
3行目でモデルのパラメータを共有メモリ上に移動させている。
5~8行目でそれぞれのプロセスが異なる処理を行う。
9~10行でそれぞれのプロセスの終了。
この並列化を適用したものが、"hogwild(SGDの並列化計算)"でpytorch/example/にコードが上がっているので参照するとイメージがつかめるはず。
自分のパソコンで作成中の強化学習モデルのある部分を切り出して、CPU並列化を適用した時としなかった時で比較して見た。 結果:20s(並列化なし), 31s x 2 (並列化有り、プロセス数:2)

GPUでの並列処理

(参照: Multi-GPU examples — PyTorch Tutorials 0.4.0 documentation)

import torch
mynet = torch.nn.DataParallel(MyModel, device_ids=[0, 1, 2])

 "nn.DataParallel(module, device_ids=None)"とすることで指定した複数gpuで(defaultでは全GPU)バッチ処理の並列処理をおこなう。そのため、指定するdevice数はバッチサイズよりは小さくないといけなく、理想としてはbatch sizeがdevice数の定数倍。
 似たようなものとして"nn.parallel.DistributedDataParallel(module, device_ids=None)"があるが、これを使うには制約が多そうなので上のnn.DataParallelを使うのが無難そう。

CPU+GPU

(Module内でCPU処理とGPU処理を混在させたとき)
(参照: Multi-GPU examples — PyTorch Tutorials 0.4.0 documentation)

import torch
import torch.nn as nn

device = torch.device("cuda:0")

class DistributedModel(nn.Module):

    def __init__(self):
        super().__init__(
            embedding=nn.Embedding(1000, 10),
            rnn=nn.Linear(10, 10).to(device),
        )

    def forward(self, x):
        # Compute embedding on CPU
        x = self.embedding(x)

        # Transfer to GPU
        x = x.to(device)

        # Compute RNN on GPU
        x = self.rnn(x)
        return x

 このクラスを呼び出した時、embedding()の部分はCPUで計算、それ以降はGPUにデータを移動させてGPUで計算をおこなう。このmodelに対してnn.DataParallel()を適用すると、gpu部分のみ並列化される(はず?)。