통상의 GPU에 의한 병렬화는 Data Parallel(데이터 병렬)이며, 데이터 배치를 복수의 GPU로 병렬로 처리하는 것입니다.
PyTorch라면 nn.DataParallel을 사용하는 방법입니다. (참고: htps : // 이 m / 어느 테마 47 / ms / 2b92f94c734b0 a 11609d )
이제 학습 속도를 빠르게 할 수 있습니다. 그러나 데이터 병렬로는 하나의 GPU 용량을 넘은 네트워크를 학습할 수 없습니다.
그래서 또 다른 병렬화 방법으로 하나의 네트워크를 여러 GPU로 부분 분할하여 학습하는 방법이 생각된다.
이것을 Model Parallel(모델 병렬)이라고 합니다.

특히 3차원 데이터의 Deep Learning 등에서는 네트워크가 거대해져 GPU에서의 메모리가 부족해지는 사태가 발생합니다.
그런 장면에서이 Model Parallel은 도움이 될 것입니다.

구현하다

이 방법은 PyTorch 튜토리얼을 참조했습니다.
htps : // py와 rch. 오 rg / 쓰리 아 ls / 어서 r 메아 아테 / 만약 l_ 파랏 ぇ l_ 쓰리 아 l. HTML
ResNet50을 병렬화합니다.
모델 병렬의 이미지로서는, 아래 그림과 같은 느낌입니다.
네트워크 전반부를 GPU1로, 후반부를 GPU2로 학습합니다.

환경

GPU는 Tesla P100(16GB)의 GPU를 2대 사용했습니다. (NVLINK로 연결됨)
OS는 Linux(Red Hat Enterprise)
CUDA 9.0
Anaconda
파이썬 3.6
PyTorch 1.1.0

코드

ResNet의 네트워크를 2개로 분할하고 각각을 GPU('cuda:0'와 'cuda:1')를 지정해 할당합니다.

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.models as models
from torchvision.models.resnet import ResNet, Bottleneck

num_classes = 1000
class ModelParallelResNet50(ResNet):
    def __init__(self, *args, **kwargs):
        super(ModelParallelResNet50, self).__init__(Bottleneck, [3, 4, 6, 3], num_classes=num_classes, *args, **kwargs)
        self.seq1 = nn.Sequential(
            self.conv1,
            self.bn1,
            self.relu,
            self.maxpool,
            self.layer1,
            self.layer2
        ).to('cuda:0')
        self.seq2 = nn.Sequential(
            self.layer3,
            self.layer4,
            self.avgpool,
        ).to('cuda:1')
        self.fc.to('cuda:1')
    def forward(self, x):
        x = self.seq2(self.seq1(x).to('cuda:1'))
        return self.fc(x.view(x.size(0), -1))

입력은 무작위 값으로 채워진(3,1024,1024) 이미지로 하고 batch size=10으로 학습해 봅니다.

batch_size = 10
image_w = 1024
image_h = 1024
model = ModelParallelResNet50()
#通常のResNet50で学習するときは　model = models.resnet50(num_classes=num_classes).to('cuda:0')
model.train()
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)
one_hot_indices = torch.LongTensor(batch_size).random_(0, num_classes).view(batch_size, 1)
for _ in range(1000):
    inputs = torch.randn(batch_size, 3, image_w, image_h)
    labels = torch.zeros(batch_size, num_classes).scatter_(1, one_hot_indices, 1)
    optimizer.zero_grad()
    outputs = model(inputs.to('cuda:0'))
    labels = labels.to(outputs.device)
    loss_fn(outputs, labels).backward()
    optimizer.step()

실행하는 동안 쉘에서 nvidia-smi 명령을 실행하여 GPU 사용량을 확인합니다.

$ nvidia-smi
...
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla P100-SXM2...  On   | 00000000:61:00.0 Off |                    0 |
| N/A   31C    P0    46W / 300W |  14637MiB / 16276MiB |     48%      Default |
+-------------------------------+----------------------+----------------------+
|   1  Tesla P100-SXM2...  On   | 00000000:62:00.0 Off |                    0 |
| N/A   35C    P0   241W / 300W |   6657MiB / 16276MiB |     80%      Default |
+-------------------------------+----------------------+----------------------+
...

2개의 GPU에서 메모리가 사용되고 있는 것을 확인할 수 있었습니다.

마지막으로

nn.Sequential로 레이어를 정리해 GPU를 지정하는 것만으로, 간단하게 할 수 있었습니다.
덧붙여서 거대한 네트워크를 학습하는 다른 방법으로는 Unified Memory 등이 있습니다.
참고 : htps : // 이 m / 이것 / ms / 4494442 에 b71 베아 0 b5b2

또 다른 방법으로는 원래의 GPU 메모리 용량이 큰 TPU를 사용하는 것도 생각할 수 있습니다.

＊추기(2020/04/30)
3D U-Net의 Model Parallel을 구현한 것을 github에 공개했습니다.
htps : // 기주 b. 코 m / 아타케 히로 / 3D - 네 - tpy와 rch - l-parap 1