Multi GPU with Pytorch (DistributedDataParallel)

1. Introduction

많은 연구자 및 개발자들이 관심을 갖는 주제 중 하나는 Deep Learning 모델을 빠르게 학습시키는 방법입니다. 일반적으로 Deep Learning 모델을 빠르게 학습시키는 방법에는 여러가지가 있습니다. 당연히 CPU가 아닌 GPU를, GPU가 아닌 TPU를 사용함으로써 학습 속도를 향상시킬 수 있습니다. 성능이 더 좋은 장비를 사용하는 것 또한 한 가지 방법이 될 수 있습니다. 당연히 GTX1050보다 GTX1080ti를 사용하는 것이 빠를 것입니다. 하지만 비용을 들여 더 좋은 장비를 구매하는 것에는 한계가 있습니다.

그래서 많은 Machine Learning Framework에서는 여러 대의 GPU를 이용하여 학습을 가속화하는 방법을 제공하고 있습니다. 이 글에서는 많은 Machine Learning Framework 중 Pytorch 내부의 DistributedDataParallel을 이용해서 Multi GPU를 통한 학습 가속화 방법을 소개합니다.

 

2. Overall Flow

먼저, 코드를 보여드리기 전에 DistributedDataParallel의 전체적인 흐름을 짚겠습니다. 아래의 흐름은 4개의 GPU를 사용하며 전체 데이터의 개수를 1000개로 가정하고 있습니다.

 

1. 먼저 사용할 GPU의 개수를 count합니다.
2. GPU의 개수만큼 Process를 생성합니다.
   4개의 GPU를 사용하기 때문에 4개의 Process가 생성될 것입니다.
3. 각 Process에 GPU를 순서대로 할당합니다.
4. Process마다 Model을 생성합니다.
   이 Model을 해당 GPU에 할당하며 학습시 Loss를 공유하도록 설정합니다.
5. Data를 Process의 개수에 맞게 분할합니다. 분할된 Data들을 각 Process에 할당합니다.
   (1000 / 4) = 250 개의 데이터들이 각 Process에 할당됩니다.
6. 각 Process에서 Batch Data를 생성합니다.
7. Model이 Batch Data에 대한 Inference를 수행하고 Loss를 계산합니다.
8. 모든 Process에서 Loss를 계산하면 backward를 수행합니다.
9. 6~8을 반복합니다.

 

위의 과정을 보면 Pytorch의 DistributedDataParallel은 Data를 Process의 개수만큼 분할하여 학습하는 것으로 생각할 수 있습니다.

3. DataLoader

import time
import torch
import torchvision
import model as baseline

from torchvision import transforms

def get_train_loader(image_size, batch_size, num_worker):
    transform_train = transforms.Compose([
        transforms.RandomResizedCrop(image_size),
        transforms.RandomHorizontalFlip(),
        transforms.RandomVerticalFlip(),
        transforms.RandomRotation(45),
        transforms.RandomAffine(45),
        transforms.ColorJitter(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])])
    train_datasets = torchvision.datasets.ImageFolder(
        root=f'data/test', transform=transform_train)
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_datasets)
    shuffle = False
    pin_memory = True
    train_loader = torch.utils.data.DataLoader(
        dataset=train_datasets, batch_size=batch_size, pin_memory=pin_memory,
        num_workers=num_worker, shuffle=shuffle, sampler=train_sampler)
    return train_loader

일반적인 DataLoader를 생성하는 방법과 다르지 않지만 21번째 줄에서 DistributedSampler 를 선언하는 것과 24~26 줄에서 sampler=train_sampler 로 선언하는 것이 중요합니다. 또한 DistributedSampler 를 사용할 경우 전체 DataLoader에서 shuffle을 하는 기능은 사용할 수 없으나 각 Process로 분할된 DataLoader내에서는 shuffle이 가능합니다.

 

DistributedSampler Shuffle 방법

4. Main

def main():
    ngpus_per_node = torch.cuda.device_count()
    world_size = ngpus_per_node

    torch.multiprocessing.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, ))

2번째 줄에서 전체 Process에서 접근할 수 있는 GPU의 개수를 ngpus_per_node 로 선언한 후 Process를 생성하고 main_worker함수를 실행 합니다.

5. Main Worker

def main_worker(gpu, ngpus_per_node):

    image_size = 224
    batch_size = 512
    num_worker = 8
    epochs = 1

    batch_size = int(batch_size / ngpus_per_node)
    num_worker = int(num_worker / ngpus_per_node)
    
    torch.distributed.init_process_group(
            backend='nccl',
            init_method='tcp://127.0.0.1:3456',
            world_size=ngpus_per_node,
            rank=gpu)
    model = baseline.ResnetModel()
    torch.cuda.set_device(gpu)
    model = model.cuda(gpu)
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[gpu])

    train_loader = get_train_loader(
        image_size=image_size,
        batch_size=batch_size,
        num_worker=num_worker)

    optimizer = torch.optim.SGD(
        params=model.parameters(),
        lr=0.001,
        momentum=0.9)
    criterion = torch.nn.CrossEntropyLoss().to(gpu)

    model.train()
    for epoch in range(epochs):

        start_time = time.time()
        for j, (images, labels) in enumerate(train_loader):
            images, labels = images.to(gpu), labels.to(gpu)
            
            optimizer.zero_grad()
            logits, _, _ = model(images)
            loss = criterion(logits, labels)
            loss.backward()
            optimizer.step()


            print(f'epoch : {epoch} | step : {j} / {len(train_loader)} | mp : {gpu}')
        end_time = time.time()
        print('total time :', end_time - start_time)

4번째 줄에서 한 step에 학습할 Batch 수를 설정합니다. 8번째 줄에서 각 프로세스가 한 step에 학습할 Batch 수를 설정합니다. 11~15번째 줄에서 어떤 ip와 port를 통해서 Loss를 공유할지 설정합니다. 17번째 줄에서 각 Process에 GPU를 순서대로 할당합니다.

이후로는 기본적인 pytorch 학습 코드와 같으므로 생략하겠습니다.

 

6. Test

해당 코드를 실행할때 GPU를 1, 2, 4인 상황에 대해서 1 epoch를 학습하는데 소요되는 시간을 비교해보았습니다.

+-----+------------+-------------------+----------+--------+
| GPU | step batch | batch per process | gpu util |  time  |
+-----+------------+-------------------+----------+--------+
|  1  |    128     |        128        |    99 %  | 128.84 | 
|  2  |    256     |        128        |    99 %  |  66.5  |
|  4  |    512     |        128        |    99 %  |  37.8  |
+-----+------------+-------------------+----------+--------+

 

4개의 GPU를 사용할 때의 GPU-Util

7. Comment

이 글을 통해서 많은 연구자 분들이 효율적으로 장비를 활용할 수 있기를 기원합니다.

[##차금강##]