npz compressed를 이용한 데이터 로드 속도 개선 (Feat. 학습 속도를 부스팅하기 위한 다양한 시도들)

목차

1. 학습 속도가 너무 느려요... GPU Util이 안나와요...😭

2. npz를 npz compressed로 바꿔서 저장해보자!

    2-1. npz를 불러와서 npz compressed로 저장하는 코드

    2-2. npz vs npz compressed 로딩 속도 비교

3. 학습 속도 개선을 위해 시도해 볼 만한 방법들

---

1. 학습 속도가 너무 느려요... GPU Util이 안나와요...😭

현재 연구중인 베이스라인 모델이 DDP (Distributed Data Parallel)를 사용하고 있음에도 불구하고, training 속도가 느린 문제가 있었습니다. 학습 속도를 개선하기 위해 여러 가지 방법을 시도해보았지만, GPU Util이 제대로 나오지도 않고 epoch 하나 도는데 무려 10시간이 걸려 멘탈이 탈탈 털릴 때쯤.. 될대로 돼라 하고 시도했는데 크게 도움을 받은 아주 간단한 방법을 소개드리려고 합니다.

 

2. npz를 npz compressed로 바꿔서 저장해보자!

제 모델의 경우 30만 개가 넘는 wav 파일과 npz 파일을 처리해야 하기 때문에, 데이터 로딩 과정에서 병목 현상이 발생했을 가능성이 높았습니다. 아래 블로그 글에서 npz compressed가 npz보다 로딩 속도가 훨씬 빠르다고 해서 한 번 시도해보았습니다. 

 

npz vs pickle

 

결론부터 말씀드리자면, npz에서 npz compressed로 파일 확장자 하나 바뀌는 것일 뿐인데 허무할 정도로 로딩 속도가 빨랐습니다. 저장 속도가 느리다곤 하지만 체감상 그렇게 느리진 않은 것 같습니다. 문제가 해결돼서 기쁘지만 그동안 뭔가 손해본 느낌...🙄

 

2-1. npz를 불러와서 npz compressed로 저장하는 코드

import os
import numpy as np
import glob
from tqdm import tqdm

input_dir = "./data/dev/"
output_dir = "./data/dev_cprd/"

os.makedirs(output_dir, exist_ok=True)

files = [file for file in glob.glob(input_dir+'*.npz')]
for file in tqdm(files, desc="Compressing files", unit="file"):
    data = np.load(file, allow_pickle=True)
    file_name = os.path.basename(file)
    output_file = os.path.join(output_dir, file_name)
    np.savez_compressed(output_file, data)

 

2-2. npz  vs npz compressed 로딩 속도 비교

import os
import numpy as np
import glob
from tqdm import tqdm
import time

npz_path = [f"./data/train/{i}.npz" for i in range(1000)]
npz_cprd_path = [f"./data/train_cprd/{i}.npz" for i in range(1000)]

start_time = time.time()
for file in tqdm(npz_path, desc="npz load", unit="file"):
    data_npz = np.load(file, allow_pickle=True)
end_time = time.time()
npz_time = end_time - start_time

start_time = time.time()
for file in tqdm(npz_cprd_path, desc="npz compressed load", unit="file"):
    data_npz_cprd = np.load(file, allow_pickle=True)
end_time = time.time()
npz_cprd_time = end_time - start_time

print(f"npz 파일: {npz_time:.4f} 초")
print(f"npz compressed 파일: {npz_cprd_time:.4f} 초")

 

호오....

 

속이 다 시원..

 

3. 학습 속도 개선을 위해 시도해 볼 만한 방법들

1. 데이터셋 attribute 중 학습에 필요없는 거 삭제하기

2. num_workers 수(=사용할 cpu 스레드 수) 늘리기: 적당한 값으로 설정해야함. 너무 작으면 multi processing 효과를 보기 힘듦. 너무 많으면 오버헤드가 커져 오히려 느려짐.

3. pin_memory=True: CPU 메모리에 있는 데이터를 GPU로 더 빠르게 전송함. 

4. prefetch_factor 늘리기: 각 worker가 한 번에 가져오는 배치 수를 결정하기 때문에 값을 늘리면 더 많은 데이터를 미리 가져와서 CPU가 처리하는 작업이 늘어나게 됨.

5. persistent_workers=True: 데이터 로더의 모든 배치를 처리한 후에도 worker가 종료되지 않고 유지됨. 데이터 로드 시 오버헤드 줄이는 효과를 기대할 수 있음.