[ PyTorch ] DataLoader 잘 짜서 병목 극복하기, GPU util 높이기

 

문제 상황

GPU util이 왜 항상 20% 정도만 나올까. 정말 모델이 작아서 그런걸까?

OOM이 뜨기 직전의 상황에도 왜 util은 100%를 찍지 않는 걸까.

이런 고민들을 해결해 보려고 초점을 맞춰 보겠다.

원인

우선은 조사한 결과 가장 많이 영향을 주는 것은

데이터를 load 하는 과정과 feed 하는 과정 사이에서 생기는 딜레이 때문이다.

해외의 pytorch 포럼이나 nvidia 포럼을 보아도 동일한 답변이고

스텍오버플로우나 기타 블로그에서도 데이터 로더의 설계문제라는 사실을 많이 지적한다.

따라서 데이터 로더를 좀 손보는 방향으로 실험을 진행했다.

 

 

데이터 로더 실험

이미지를 불러올때 몇가지 trade-off와 싸워야 한다.

흔히 말하는 공간복잡도와 시간복잡도의 싸움인데

메모리가 허락하는 한 메모리에 모든 데이터를 올려보는 것이 좋다.

하지만 메모리에 올리는 동안 시간이 얼마나 걸릴지 알 수없기 때문에

이득인지 손해인지 알려면 일단 한번 해 봐야 한다.

실험의 설계는 간단하다.

우선 데이터의 위치만 가지고 있고 실시간으로 필요한 양만큼만 불러오는 과정이다.

torch 데이터 로더에서는 __init__()에 구현하느냐 __getitem__()에 구현하느냐의 차이가 되겠다.

실험 1

데이터를 그때그때 필요한 양만큼만 불러오자.

class ImageDataset(Dataset):
    def __init__(self, data_path, pre_transforms):
        df = pd.read_csv(
            data_path + "datapath_with_label.csv"
        ) # 파일의 경로, 라벨만 저장되어있는 csv

        df["file"] = df["file"].map(lambda x : "../input/" + x)

        self.X = df["file"] # 파일의 경로
        self.y = df["label"] # 라벨

        self.transforms = pre_transforms
    def __len__(self):
        len_dataset = len(self.X)
        return len_dataset

    def __getitem__(self, idx):
        X, y = self.transforms(Image.open(self.X[idx])), self.y[idx]
		# 여기서 이미지를 불러오고 전처리까지 함
        return X, torch.tensor(y)

첫번째는 이런식으로 DataFrame에 파일의 이름과 라벨만 넣어두고 getitem을 할때 PIL의 Image 로 불러오게 된다.
원본 이미지의 크기가 512 * 384 * 3 인 상황인데

tic = time.time()
dataset = ImageDataset(
    data_path = "../input/data/train/",
    pre_transforms = transforms.Compose([
        transforms.Resize((512//4,384//4)),
        transforms.CenterCrop((64, 64)),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.ToTensor()
    ])
)
print(time.time() - tic)

이렇게 1/4로 축소하고 크롭한 뒤 불러오는 작업을 수행한다.
이렇게 dataset을 만드는데 걸린 시간은 0.026030540466308594초로 굉장히 빠르게 수행된다.

dataloader = DataLoader(
    dataset,
    batch_size  = 100,
    shuffle     = True,
    sampler     = None,
    num_workers = 1
)

그리고 이렇게 100개의 배치를 뽑아주는 데이터 로더를 만들어서 간단한 Convolution 하나를 연산해 보자

device = torch.device("cuda:0")
tictoc, iteration = 0, 10
for _ in range(iteration):
    tic = time.time()
    for X, y in iter(dataloader):
        single_batch = torch.nn.Conv2d(3,3,(3,3),device = device)(X.to(device))
    tictoc += time.time() - tic
print(tictoc / iteration)

10 epoch 정도 반복한 후에 평균적인 시간을 구해보면

 

평균이 89초... 그만 알아보자.

 

 

실험 2

이번에는 이미지를 일단 다 불러와서 init 할때 메모리에 왕창 올려두고 연산을 해보자.

우선은 이미지 한장당 0.5MB 정도 하기 때문에 꽤 큰 편에 속한다.

그래서 transfrom 과정을 적절하게 나누어서 Flip같은 과정은 getitem 할때마다 사용되도록 나눠보겠다.

 

class ImageDataset(Dataset):
    def __init__(self, data_path, pre_transforms, transforms):
        df = pd.read_csv(
            data_path + "datapath_with_label.csv"
        )
        
        df["file"] = df["file"].map(lambda x : "../input/" + x)
        self.X = []
        for X in df["file"]:
            self.X.append(pre_transforms(Image.open(X)))
            # 여기서 이미지를 불러오고 크기를 줄이는 전처리만 해서 리스트에 담는다
        self.y = df["label"]
        
        self.transforms = transforms
    def __len__(self):
        len_dataset = len(self.X)
        return len_dataset
    
    def __getitem__(self, idx):
        X, y = self.transforms(self.X[idx]), self.y[idx]
        # 뒤집기등의 augmentation은 여기서 해야 불러올때마다 적용된다.
        
        return X, torch.tensor(y)

 

우선은 위와같이 2개의 transform을 받도록 해보았다.

이미지를 자르고 해상도를 바꾸는 과정은 init에 나머지 과정을 getitem에 넣었다.

tic = time.time()
dataset = ImageDataset(
    data_path = "../input/data/train/",
    pre_transforms = transforms.Compose([
        transforms.Resize((512//3,384//3)),
        transforms.CenterCrop((64, 64)),
    ]),
    transforms = transforms.Compose([
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.ToTensor()
    ])
)
print(time.time() - tic)

이렇게 두개로 나누어진 같은 작업을 올려주면

 

81초. 데이터세트 자체가 생성되는데는 크게 시간이 늘었다.

그리고 똑같이 100개의 배치를 만들어주는 데이터 로더에 물려보면

 

1에폭당 평균 7초가 걸린다. 10%도 안되는 짧은 시간이다.

nvidia-smi로 확인해 본 결과 실험 1에서는 gpu util 이 0% 로 나오는 반면

실험 2는 3~5% 로 활발하게 GPU 가 돌아가는 모습을 보였다.

 

 

 

결론

gpu util은 초당 해낼 수 있는 연산량에 비해 초당 수행되는 연산량을 정량화한 것이라고 한다.

결국 gpu메모리에 올려야하는 자료를 얼마나 가깝게 준비하고 있느냐가 gpu util에 큰 영향을 줄 수 있다.

따라서 최대한 메모리에 올릴 수 있는 만큼 올려서 gpu메모리와 가까운 곳에 자료를 위치시키는 것이 좋겠다.

 

//참고

https://ainote.tistory.com/14

GPU Util 99% 달성하기