[NVIDIA] DALI multi-GPU 사용법 with PyTorch

※ 해당 글은 vision ai, classification task 관련된 내용만 다룹니다.

오늘은 이전 글에 이어서 DALI dataloader를 multi-gpu로 load하는 방법을 설명드리고 single-gpu와 multi-gpu간의 속도 차이를 확인해보겠습니다. 그리고 이전 글에서 추가된 부분만 설명드리도록 하겠습니다. 

0.  Experiment Setup

DALI를 적용해볼 data type은 image와 video 이며 실험환경은 다음과 같습니다.

• PyPI
  • PyTorch: 1.10.0
  • DALI: 1.6
  • decord: 0.6.0
  • OpenCV: 4.5.3.56
• Hardware
  • CPU: Intel(R) Xeon(R) Gold 5120 CPU @ 2.20GHz 
  • GPU: Tesla V100-PCIE-32G
• Dataset
  • Image: 21,453 VOC images 
  • Video: 173 random sampled videos
    • Info: 300 frames, 1080p

아래 설명에 사용한 모든 코드는 여기에 올려두었습니다.

1.  Multi-GPU DALI Image Loader

# multigpu_dali_imageloader.py
import os
import argparse
import time

import torch
import torch.distributed as dist

from nvidia.dali import pipeline_def
from nvidia.dali.plugin import pytorch
import nvidia.dali.fn as fn
import nvidia.dali.types as types

LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) 
RANK = int(os.getenv('RANK', -1))
WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))

@pipeline_def
def image_pipe(file_root: str,
               local_rank: int,
               world_size: int,
               image_size: int=640):
    jpegs, labels = fn.readers.file(file_root=file_root,
                                    initial_fill=1024,
                                    random_shuffle=True,
                                    shard_id=local_rank, # added for multi-gpu
                                    num_shards=world_size, # added for multi-gpu
                                    name="Reader")
    images = fn.decoders.image(jpegs, 
                               device="mixed", 
                               output_type=types.RGB)
    
    images = fn.resize(images, 
                       device="gpu", 
                       size=[image_size, image_size],
                       interp_type=types.INTERP_LINEAR)
    return images, labels[0]

class DALIImageLoader():
    def __init__(self, 
                 path: str, 
                 batch_size: int, 
                 num_threads: int,
                 local_rank: int,
                 world_size: int):
        pipe = image_pipe(batch_size=batch_size,
                          num_threads=num_threads, 
                          device_id=local_rank, # added for multi-gpu
                          local_rank=local_rank,
                          world_size=world_size,
                          file_root=path,
                          seed=123456)
        pipe.build()

        self.dali_iterator = pytorch.DALIGenericIterator(pipe,
                                                         ["data", "label"],
                                                         reader_name="Reader",
                                                         last_batch_policy=pytorch.LastBatchPolicy.PARTIAL,
                                                         auto_reset=True)
    def __len__(self):
        return int(self.epoch_size)

    def __iter__(self):
        return self.dali_iterator.__iter__()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify')
    parser.add_argument('--device', default='6,7', help='cuda device, i.e. 0 or 0,1,2,3')
    parser.add_argument('--batch_size', type=int, default=32, help='total batch size for all GPUs, -1 for autobatch')
    parser.add_argument('--num_threads', type=int, default=8, help='number of threads')
    parser.add_argument('--data_dir', type=str, default='/usr/src/app/da2so/datasets/VOC/images', help='dataset directory')    
    args = parser.parse_args()

    os.environ['CUDA_VISIBLE_DEVICES'] = args.device 
    torch.cuda.set_device(LOCAL_RANK)
    device = torch.device('cuda', LOCAL_RANK)
    os.environ['NCCL_BLOCKING_WAIT'] = '1'  # set to enforce timeout
    dist.init_process_group('nccl' if dist.is_nccl_available() else 'gloo')

    daliloader = DALIImageLoader(path=args.data_dir,
                                 batch_size=args.batch_size,
                                 num_threads=args.num_threads,
                                 local_rank=LOCAL_RANK,
                                 world_size=WORLD_SIZE)
    if RANK == 0:
        start_time = time.time()
        for idx, inp in enumerate(daliloader):
            print(f'image shape: {inp[0]["data"].shape}')
            print(f'label shape: {inp[0]["label"].shape}')
        print(f'[Multi-GPU {args.device} DALI Imageloader] time: {time.time() - start_time}')

Multi-GPU를 사용하기 위해 기존 DALIImageLoader에서 추가된 arugments는 local_rank, world_size 2개입니다. Multi-GPU학습에 익숙하신분들은 아시겠지만 local_rank는 사용되는 GPU number(id)를 뜻하고 world_size는 사용되는 GPU 전체 개수를 의미합니다. 해당 인자들은 DALIImageLoader의 image_pipe함수에 사용됩니다.

• image_pipe
  • device_id=local_rank: local_rank에 해당하는 하나의 GPU를 할당
    • 해당 인자는 @pipeline_def decorator의 파라미터임

device_id=local_rank을 통해 각 GPU로 pipeline을 실행할 수 있게 되었습니다. 여기서 추가적으로 각 GPU가 서로 다른 samples을 managing할 수 있도록 하는 기술인 sharding을 사용합니다. Dataset을 여러 parts(shards)로 나누어 각 GPU는 고유의 shard로 data load를 진행하게 됩니다.

Sharding

• fn.readers.file
  • shard_id=local_rank: 각 GPU가 고유의 shard_id를 가지도록 함
  • num_shards=world_size: 사용하는 총 shard개수를 지정

이제 위 코드를 torchrun --standalone --nnodes=1 --nproc_per_node=2 multigpu_dali_imageloader.py --device=0,1 명령어로 실행시켜봅니다. GPU 2개를 사용하여 각 GPU당 32 batch size으로 설정하였고 image는 640으로 resize 하였습니다.

총 21,453개의 image를 load하는데 3.1초 정도밖에 걸리지 않네요! 

2. Multi-GPU DALI Video Loader

Video loader부분에 대한 설명은 위의 image loader부분과 내용이 다수 겹치고 기존 single-gpu video loader부분과도 유사한 부분이 많아 바뀐 코드만 보여드립니다. 

... 생략

@pipeline_def
def video_pipe(filenames: List[str],
               labels: List[int], 
               sequence_length: int,
               stride: int,
               local_rank: int,
               world_size: int):

    videos, label = fn.readers.video(device="gpu", 
                              filenames=filenames,
                              labels=labels, 
                              sequence_length=sequence_length,
                              normalized=False, 
                              random_shuffle=True, 
                              image_type=types.RGB,
                              dtype=types.UINT8,
                              initial_fill=16,
                              num_shards=world_size, # added for multi-gpu
                              shard_id=local_rank, # added for multi-gpu
                              stride=stride,
                              name="Reader")
    return videos, label[0]

... 생략

image loader와 다른점은 fn.readers.video 함수의 인자로 shard_id와 num_shards를 인자를 추가해주어야 한다는 점입니다.

Multi-GPU DALI video loader에서도 위 코드를 torchrun --standalone --nnodes=1 --nproc_per_node=2 multigpu_dali_videoloader.py --device=0,1  명령어로 실행시켜봅니다. GPU 2개로 각 GPU당 8  batch size를 가지도록 하였습니다.

총 173개의 video를 load하는데 48.2초 소요되었습니다.

 

2. Single-GPU vs Multi-GPU DALI Loader 시간 비교

해당 섹션에서는 GPU개수에 따른 DALI Loader의 시간 측정을 진행합니다. 추가적으로 Pytorch dataloader의 data load시간을 baseline으로 잡고 진행하였습니다.

• Experiment Setting
  • [DALI, PyTorch] batch size = 32(image), 8(video)
    
    • 각 GPU당 batch size임
  • [DALI] num threads = 8
  • [PyTorch] num workers = 8
  • [PyTorch] pin_memory = True

	Time cost for image loader (s)	Time cost for video loader (s)
PyTorch CPU	101.06	59.01
DALI 1-GPU	6.24	97.45
DALI 2-GPU	3.07	48.28
DALI 4-GPU	1.57	21.34

• Image loader 결과
  • CPU로 data load하는 PyTorch보다 DALI사용 시 data load속도가 훨씬 빠름
  • GPU가 배로 늘어날수록 속도도 비례하여 빨라짐
• Video loader 결과
  • CPU로 data load하는 PyTorch보다 Single GPU를 사용하는 DALI가 더 느림
  • 하지만, GPU 2개이상 사용시 PyTorch data loader보다 빨라짐
  • Image loader와 동일하게 GPU 개수에 따라 속도도 비례하게 빨라짐