[NVIDIA] DALI 사용법 with PyTorch

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

오늘은 DALI를 사용하는 방법을 알아보고 DALI를 사용했을 때와 사용하지 않았을 경우의 time cost차이를 직접 측정해보도록 하겠습니다.

1. DALI 란?

• Data Loading Library (DALI)는 DNN을 training할때 data loading 및 pre-processing을 GPU을 사용할 수 있도록 하는 GPU-accelerated library
  • 그래서 CPU를 사용할 때보다 훨씬 빠르게 training이 가능
• DALI는 자체적인 execution engine을 가지며 input pipeline의 throughput을 최대화 시키기위해 설계됨 
• DALI는 portable하기 때문에 PyTorch, TensorFlow, MXNet에 쉽게 integrated 가능
• 다양한 data format 지원: TFRecord, COCO, JPEG, JPEG 2000, WAV, FLAC, OGG, H.264, VP9 and HEVC
• 여러 GPUs에 scaleable가능

2. DALI 사용법

2.1 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

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

2.2 DALI Image data loading

DALI를 이용하여 Image data를 loading하는 방법을 소개드립니다. 

import time

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

@pipeline_def
def image_pipe(file_root: str,
               image_size: int=640):
    jpegs, labels = fn.readers.file(file_root=file_root,
                                    initial_fill=1024,
                                    random_shuffle=True,
                                    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):
        pipe = image_pipe(batch_size=batch_size,
                          num_threads=num_threads, 
                          device_id=0, 
                          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__":
    start_time = time.time()
    daliloader = DALIImageLoader(path='/usr/src/app/da2so/datasets/VOC/images',
                                 batch_size=32,
                                 num_threads=8)
    for inp in daliloader:
        print(f'image shape: {inp[0]["data"].shape}')
        print(f'label shape: {inp[0]["label"].shape}')
    print(f'[DALI Imageloader] time: {time.time() - start_time}')

위 코드에서는 batch_size를 32로 고정하여 DALIImageLoader를 intialization하는 것을 시작으로 daliloader를 통해 data load하는 데까지 소요되는 시간을 측정하는 코드입니다. 이제 그럼 DALIImageLoader에 대해 하나하나 자세히 살펴보죠.

• DALIImageLoader
  • image_pipe 함수 호출을 통해 DALI의 pipeline object을 initalization
  • pipe.build() 를 통해 pipeline build
  • torch.DALIGenericIterator는 build한 pipeline을 대상으로 Pytorch용 DALI iterator 생성
    • 즉, DALI pipeline으로 load된 data가 gpu device로 mapping된 tensor를 형태를 가지게 됨
    • last_batch_policy: data수가 정확히 batch로 나눠지지 않을경우 마지막 batch를 어떻게 처리할 지 정함
      • PARTIAL은 마지막 batch가 setting한 batch 수보다 작을 경우에 그대로 작은 batch로 data load해줌
      • 예를 들어 전체 data가 10개이고 batch가 4라면 마지막 iteration에서는 batch가 2로 설정
    • auto_reset=True: DALI iterator의 마지막 iteration에서 StopIteration이 발생하고 자동적으로 reset() 호출해줌
• image_pipe
  • @pipeline_def라는 decorator를 사용해서 DALI pipeline구성하게 해줌
    
    • 해당 decorator를 arugment로 image_pipe함수의 batch_size, num_threads, device_id
    • device_id: gpu device id를 의미
  • fn.reader.file: file 또는 directory경로를 입력으로 image content와 label을 return (아래 그림 참조)
    • 해당 process는 cpu로 동작함
    • file_root: data files을 담고있는 directory
    • initial_fill: shuffling에 사용될 buffer_size
    • random_shuffle: data shuffling 유무
    • name: torch.DALIGenericIterator의 reader_name과 일치되는 이름
  • fn.decoders.image: image content를 decoding하는 processing
    • device='mixed': cpu와 gpu를 mix해서 사용
      • jpeg인 경우 nvJPEG library(or libjpeg-turbo)을 사용하고 다른 image format일 경우 OpenCV사용
    • fn.resize: batch로 data load하기위해 image size를 동일시 함
      • device='gpu'은 gpu를 통해 resize operation진행함

fn.reader.file return 형식

위 코드 실행하면 아래와 같이 정상적으로 실행되는 것을 확인가능합니다. 총 21,453개의 image를 load하는데 6.2초 정도밖에 걸리지 않네요!

아래사진을 통해서는 DALIImageLoader를 실행시킴으로써 GPU memory를 사용하는 것을 확인가능합니다.

 

2.3 DALI Video data loading

이번에는 DALI를 이용하여 video를 load하는 방법을 소개드립니다. Image를 load하는 부분과 동일한 부분은 생략하고 설명드리도록 하겠습니다.

import os
import time
import glob
from typing import List
from pathlib import Path

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

VID_FORMATS = 'avi', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg' # include video suffixes

@pipeline_def
def video_pipe(filenames: List[str],
               labels: List[int], 
               sequence_length: int,
               stride: 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,
                              stride=stride,
                              name="Reader")
    return videos, label[0]

def video2label_paths(video_path: List) -> List:        
    return [int(Path(x).parts[-2]) for x in video_path]


class DALIVideoLoader():
    def __init__(self, 
                 path: str, 
                 batch_size: int, 
                 num_threads: int,
                 sequence_length: int,
                 stride: int):
        
        try:
            f = [] # video files 
            for p in path if isinstance(path, list) else [path]:
                p = Path(p) 
                if p.is_dir():  # dir
                    f += glob.glob(str(p / '**' / '*.*'), recursive=True)
                elif p.is_file():  # file
                    with open(p) as t:
                        t = t.read().strip().splitlines()
                        parent = str(p.parent) + os.sep
                        f += [x.replace('./', parent) if x.startswith('./') else x for x in t]  # local to global path
        except Exception as e:
            raise Exception(f'Error loading data from {path}: {e}\n')
        self.vd_files = sorted(x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in VID_FORMATS)
        assert self.vd_files, f'No videos found'
        self.labels = video2label_paths(self.vd_files)
        assert len(self.vd_files) == len(self.labels), f'The number of video files are not matched with label files'

        pipe = video_pipe(batch_size=batch_size, 
                          num_threads=num_threads, 
                          device_id=0, 
                          filenames=self.vd_files,
                          labels=self.labels,
                          stride=stride,
                          sequence_length=sequence_length,
                          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__":
    start_time = time.time()
    daliloader = DALIVideoLoader(path='./videos',
                                 sequence_length=60,
                                 stride=5,
                                 batch_size=8,
                                 num_threads=8)
    for inp in daliloader:
        print(f'video shape: {inp[0]["data"].shape}')
        print(f'label shape: {inp[0]["label"].shape}')
    print(f'[DALI Videoloader] time: {time.time() - start_time}')

Image파트와 비슷하게 DALIVideoLoader를 intialization하고 data load하는 데까지 소요되는 시간을 측정하는 코드입니다. 여기서 DALIVideoLoader의 __init__함수에서 image_pipe 아닌 video_pipe를 호출하는 것과 path를 입력으로 video files과 그에 대한 labels를 추출하는 것 말고는 DALIImageLoader와 크게 다르지 않습니다.

• DALIVideoLoader
  • self.vd_files: path directory에 포함된 모든 video files path
  • self.labels: video파일에 대응되는 labels
• video_pipe
  • sequence_length: video에서 가져올 frame 수
  • stride: 가져올 frame간의 interval
  • normalized: video영상을 normalize 할건지
  • image_type: video의 각 frame(image)의 type을 명시

위의 sequence_length가 60, stride가 5라는 것은 video로부터 1, 5, 10, 15, ...,295, 300 번째 frames(총 60 frames)을 data load하겠다는 의미입니다. 위 파일을 실행시키면 아래 사진과 같이 정상적으로 실행이 됨을 확인가능합니다. 하지만 video가 173개 load하는 데 97초나 걸리네요...  (많이 느리네요..ㅠ)

Video load할 경우에 GPU memory를 사용하긴 하는데 GPU utils이 image loader랑 다르게 너무 낮은 느낌이... 있네요. 

 

2.4 DALI vs OpenCV(decord) 속도 비교

해당 section에서는 DALI를 이용한 GPU data load방식과 CPU data load방식을 사용하였을때의 data load time cost를 비교하겠습니다. 비교를 위해서 GPU 1개만 사용하였고 CPU는 모두 사용하였습니다. CPU image data load는 OpenCV를, CPU video data load는 decord(cpu version)를 사용하였습니다. 

2.4.1 DALI  vs OpenCV

Tim cost비교를 위해 21,453 VOC images를 사용하였고 OpenCV를 이용한 Dataloader는 기본적인 torch.utils.data.dataloader.DataLoader를 사용하였습니다. (image size는 640으로 resize함)

위 결과로부터 DALI ImageLoader가 약 16배 정도 빠른것을 알 수 있습니다.

2.4.2 DALI  vs  decord

Time cost비교를 위해 173개의 video를 사용하였고 image부분과 동일하게 decord를 이용한 Datloader는 torch.utils.data.dataloader.DataLoader를 사용하였습니다.

위 결과를 보시면 GPU를 사용한 DALI가 약 2배정도 더 느린것을 알 수 있습니다. DALI를 이용하여 GPU 1개로 video load하는 것은 cpu보다 느리네요... 하지만, training시에 CPU data loader를 사용하였을 경우 CPU → GPU → CPU로 context switching이 많이 일어나기 때문에 이 부분에 대한 고려를 했을 때 비슷해지지않을까 싶습니다.

 

다음 글에서는 multi-gpu를 사용한 DALI에 대해 설명드리겠습니다.