[Pytorch] DistributedDataParallel(DDP)

최근 딥러닝 모델들이 더욱 커지고 복잡해지면서 단일 GPU로는 학습이 어렵거나 시간이 오래 걸리게 되었기 때문에, 이를 해결하기 위한 효과적인 방법 중 하나가 PyTorch의 DistributedDataParallel(DDP)

 

DDP(DistributedDataParallel)란?

PyTorch의 DistributedDataParallel은 여러 GPU나 여러 컴퓨터(노드)에 학습을 분산시켜 병렬로 학습 속도를 높이는 방법이며, 모델의 파라미터를 동기화하면서 데이터 배치를 분산 처리하여 효율적으로 학습할 수 있게 도와줌.

 

(1) 초기 설정

import os
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
import torch.multiprocessing as mp

def setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12360"
    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()
    
if __name__ == "__main__":
    world_size = torch.cuda.device_count()
    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)

• Rank: 프로세스(작업)의 ID 일반적으로 GPU 하나당 하나의 rank가 부여됨.
• world_size: 사용할 GPU 개수

DDP의 설정을 초기화 하여 세팅하고 종료하는 단계

 

(2) 모델 정의

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.adaptive_pool = nn.AdaptiveAvgPool2d((8, 8))
        self.fc1 = nn.Linear(128 * 8 * 8, 256)
        self.fc2 = nn.Linear(256, 100)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        x = self.adaptive_pool(x)
        x = torch.flatten(x, start_dim=1)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = MyModel().to(device)
model = DDP(model, device_ids=[rank])

 

DDP에서는 모델을 정의한 후 각 GPU마다 복사본을 만들어 사용함. 그렇기 때문에 원본 모델을 받아 각 GPU에서 병렬 학습 가능한 형태로 변환해야함.

 

(3) Sampler

train_dataset = torchvision.datasets.CIFAR100(...)
train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank)
train_loader = DataLoader(..., sampler=train_sampler)

데이터가 각 GPU에 고르게 나누어지도록 DistributedSampler를 사용하여 각 rank에 균등하게 데이터를 분배함.

예제 데이터로는 CIFAR100 사용.

 

(4) Train

학습을 시작하기 전에 epoch마다 train_sampler.set_epoch(epoch)를 호출하여 데이터를 섞어줌.

  criterion = nn.CrossEntropyLoss()
  optimizer = optim.Adam(model.parameters(), lr=0.001)
  model.train()

  for epoch in range(epochs):
      train_sampler.set_epoch(epoch)
      total_loss = 0
      
      for step, (images, labels) in enumerate(train_loader):
          images, labels = images.to(device), labels.to(device)
          optimizer.zero_grad()
          outputs = model(images)
          loss = criterion(outputs, labels)
          loss.backward()
          optimizer.step()
          
          total_loss += loss.item()
  
  cleanup()

 

 

이처럼 PyTorch의 DDP를 통해 간단하게 분산 학습을 사용할 수 있다.