[Pytorch] Pytorch DDP 사용해 multi-gpu 로 모델 학습시키기.

Pytorch는 딥러닝을 위한 프레임워크로, 많은 사람들이 사용하고 있습니다. 그 중에서도 분산 학습을 위한 DDP(Distributed Data Parallel)는 많은 연구자들이 사용하고 있는데요. 오늘은 Pytorch DDP를 사용해 모델을 학습시키는 방법에 대해 알아보도록 하겠습니다.

 

Pytorch DDP 모듈은 DP (Data Parallel)와 다르게 multi-node 에서도 사용이 가능하지만, 편의를 위해 하나의 노드, 여러개의 GPU를 가정하겠습니다.

DDP는 무엇인가요?
DDP는 분산 학습을 위한 방법 중 하나입니다. 분산 학습이란 여러 대의 컴퓨터를 사용해 모델을 학습시키는 방법으로, 빠른 학습이 가능합니다. DDP는 Pytorch에서 제공하는 분산 학습 방법 중 하나로, 여러 대의 GPU를 사용해 모델을 학습시키는 방법입니다. 이를 통해 더 빠른 학습이 가능합니다.

GPU 는 여러장 있습니다. 어떻게 사용하나요?

1. Distributed backend initialize 시키기.

import torch.distributed as dist

local_rank = int(os.environ["LOCAL_RANK"])
dist.init_process_group("nccl", init_method="env://")
args.local_rank = local_rank
torch.cuda.set_device(args.local_rank)

local_rank 는 각 GPU 의 id를 의미합니다. LOCAL_RANK 환경 변수는 torch.distributed 를 통해 자동으로 정의가 되어 있습니다.

 

2. Model wrap 시키기.

import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        return x

model = Net()
model = nn.parallel.DistributedDataParallel(model, device_ids=[rank])

임의의 Pytorch 모델을 DDP 모듈을 사용해 감싸줍시다.

 

3. Dataloader 감싸주기

# Model
train_dataset = torchvision.datasets.MNIST(
    "/MNIST/",
    train=True,
    download=True,
    transform=torchvision.transforms.Compose(
        [
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize((0.1307,), (0.3081,)),
        ]
    ),
)
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=batch_size,
    sampler=train_sampler,
)

DDP를 사용하게 되면, 각 GPU 에 데이터를 나누어 뿌려줄 친구가 필요합니다. Single GPU 로 학습시킬 때와 동일한 Dataset, Dataloader를 준비합니다. 대신, DistributedSampler를 정의한 후, DataLoader의 sampler 로 넘겨주기만 하면, 자동으로 데이터가 나뉘여 각 GPU 에 뿌려지게 됩니다.

 

4. Command

만약 GPU 를 4개 사용하고 싶다면, 다음과 같은 커맨드를 사용하면 됩니다. 저절로 LOCAL_RANK 환경변수를 세팅해주면서 (0, 1, 2, 3) multi-process를 돌려주게 됩니다.

python -m torch.distributed.launch --nproc_per_node 4 main.py

 

5. 유의할 점

이 3가지만 지키면, 기존의 single GPU 로 학습하던 모델을 multi GPU 에서 학습시킬 수 있습니다. 한 가지 유의해야할 점은 코드가 multi-process 로 돌아간다는 점입니다. 예를 들어, print('train start') 를 하게 되면, GPU 갯수만큼 print 문이 찍히게 됩니다. 때문에 로그를 찍을 때나, validation, evaluation을 할 시에는, 각 process 의 local_rank를 바탕으로 한 GPU 에서만 앞선 작업들을 시행해 주면 더 좋을 것 같습니다.

 

이만 마치겠습니다.

행복하세요~