Horovod 모듈 사용하여 Pytorch Multi GPU 학습 효율적으로 하기

1. Introduction

기본적으로 Pytorch를 이용하여 학습을 가속하기 위해서는 다수의 GPU를 활용해야 합니다. 이전 글에서 기술한 것처럼 다수의 GPU를 사용할 수 있는 방법을 Pytorch 자체에서 DistributedDataParallel이라는 모듈로 제공하고 있습니다. 하지만 해당 글에서 볼 수 있듯이 고려해야 할 사항들이 많습니다. 예를 들어 multiprocessing의 프로세스를 관리하는 것과 DataLoader에서 pin memory, shuffle 등을 고려해야 합니다. 하지만 Horovod라는 모듈을 이용하면 굉장히 간단한 코드 추가로 Multi GPU 학습이 가능합니다. 이 글에서는 Pytorch와 Horovod를 이용하여 다수의 GPU를 이용한 학습 방법을 소개합니다.

2. Installation

설치하는 방법은 굉장히 간단합니다.

pip3 install horovod # cpu installation
HOROVOD_GPU_OPERATIONS=NCCL pip3 install horovod # gpu installation    

Horovod를 사용할 때 GPU 학습을 하려면 gpu 설치 방법을 사용하고 그 외에는 cpu 설치 방법을 사용하면 됩니다.

 

3. Detail Code

import horovod.torch as hvd
hvd.init()

먼저 horovod를 사용하기 위해서는 horovod의 초기화를 진행해야합니다.

 

torch.cuda.set_device(hvd.local_rank())

다음 torch 모듈에 생성되는 프로세스에 대해 device를 할당합니다.

 

import torchvision
train_dataset = torchvision.datasets.MNIST(
    root='data',
    train=True,
    transform=torchvision.datasets.ToTensor())
train_sampler = torch.utils.data.distributed.DistributedSampler(
    trans_dataset,
    num_replicas=hvd.size(),
    rank=hvd.rank())
train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=128,
    num_workers=8,
    sampler=train_sampler)

일반적으로 사용하는 Dataset 모듈과 DataLoader 모듈을 이용해서 사용할 DataLoader를 생성합니다.
하지만 DataLoader를 생성할때 분산처리를 사용하기 때문에 DistributedSampler를 생성하여 할당합니다.
위의 Dataset 생성 코드를 해석하면 다음과 같습니다.
MNIST 데이터셋을 먼저 생성합니다. 다음 hvd.size()의 개수만큼 분할하는 Sampler를 생성합니다. MNIST 데이터셋과 Sampler를 이용하여 DataLoader를 생성합니다. 하나의 분할된 Loader는 배치사이즈가 128로 설정했습니다. 즉, GPU를 4개 생성한다면 한 스텝에 총 128 * 4개의 데이터를 학습합니다.

 

model = Model()
model = model.cuda()

Model을 생성합니다.

 

optimizer = torch.optim.Adam(
    params=model.parameters(),
    lr=1e-4)
optimizer = hvd.DistributedOptimizer(
    optimizer,
    named_parameters=model.named_parameters())

optimizer를 생성하고 horovod에서 사용할 수 있는 객체로 바꿔주기 위해 Wrapping을 합니다.

 

criterion = torch.nn.CrossEntropyLoss()
hvd.broadcast_parameters(
    model.state_dict(),
    root_rank=0)

loss function을 구성하고, 어느 프로세스를 기준으로 학습을 진행할지 설정합니다.

 

for epoch in range(100):

    acc = 0
    loss = 0
    model.train()
    for data, target in tqdm.tqdm(train_loader):

        data = data.cuda()
        target = target.cuda()

        optimizer.zero_grad()
        logits = model(data)
        step_loss = criterion(logits, target)
        step_loss.backward()
        optimizer.step()

            pred = torch.argmax(logits, axis=1)
            pred = pred.eq(target).sum().item() / data.shape[0]

            loss += step_loss.item()
            acc += pred

        print(f'loss : {loss / len(train_loader)}')
        print(f'acc : {acc / len(train_loader)}')

다음은 일반적인 학습 루프를 구성합니다.

위의 코드들을 모두 합치면 아래와 같이 구성할 수 있습니다.

import torch
import horovod.torch as hvd

hvd.init()

torch.cuda.set_device(hvd.local_rank())
train_dataset = torchvision.datasets.MNIST(
    root='data',
    train=True,
    transform=torchvision.transforms.ToTensor())
train_sampler = torch.utils.data.distributed.DistributedSampler(
    train_dataset,
    num_replicas=hvd.size(),
    rank=hvd.rank())
train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=128,
    num_workers=8,
    sampler=train_sampler)
model = Model()
model.cuda()
optimizer = torch.optim.Adam(
    params=model.parameters(),
    lr=lr)
optimizer = hvd.DistributedOptimizer(
    optimizer,
    named_parameters=model.named_parameters())
criterion = torch.nn.CrossEntropyLoss()
hvd.broadcast_parameters(
    model.state_dict(),
    root_rank=0)

for epoch in range(100):

    acc = 0
    loss = 0
    model.train()
    for data, target in tqdm.tqdm(train_loader):

        data = data.cuda()
        target = target.cuda()

        optimizer.zero_grad()
        logits = model(data)
        step_loss = criterion(logits, target)
        step_loss.backward()
        optimizer.step()

            pred = torch.argmax(logits, axis=1)
            pred = pred.eq(target).sum().item() / data.shape[0]

            loss += step_loss.item()
            acc += pred

    print(f'loss : {loss / len(train_loader)}')
    print(f'acc : {acc / len(train_loader)}')

 

4. Run

python3 train.py

일반적으로 Python 코드를 실행할 때에는 위와 같은 코드를 사용합니다. 하지만 Python에서 horovod를 함께 사용하기 위해서는 특정 커맨드를 아래와 같이 입력해야 합니다. 여기서 -np는 총 몇 개의 분산학습을 진행할지 설정하는 부분입니다.

horovodrun -np 4 python3 train.py

 

5. Result

horovodrun을 이용하여 실행하면 다음과 같이 여러 개의 GPU가 함께 사용되는 것을 확인할 수 있습니다.

[##차금강##]