문제상황

• 느린 I/O 속도때문에 병목이 생겨서 GPU를 제대로 활용하지 못함
• 큰 batch size에서, 원본 이미지의 크기가 클 경우 발생

해결방법

• 메모리 캐싱을 사용해서, 한 번 읽은 이미지는 다시 디스크를 사용하지 않고 작업

캐싱에 필요한 메모리 크기

• 만약 학습에 256x256 color image를 10만장 사용한다면, 메모리 캐싱에는 256x256x3x100000 Byte 약 19.6GB가 필요함

데이터로더에 cache 변수 만들기

class SkinDataset(Dataset):
    def __init__(self, data, cache, imsize, transform):
        self.cache = cache        
        self.imsize = imsize
        self.transform = transform
        self.data = data # self.data를 image path list로 가정

• self.cache 변수에 이미지를 저장함
• PyTorch에서 num_worker를 0보다 큰 값을 사용하면, 각각의 worker에서 만든 변수가 다음번 batch에서 초기화 되어서 캐싱을 하지 못함
• shared dict 를 사용하면 메모리 캐싱 구현 가능

    def __getitem__(self, index):
    	def imread_cache(path):
        	img = self.cache.get(path, None)
            if data is None:
            	img = cv2.imread(path)
                img = cv2.resize(img, (self.imsize, self.imsize))
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                self.cache[path] = data
            else:
	            pass
            return img
        
        path = self.data[idx]
       	img = imread_cache(path)
        img = self.transform(img)
        return img

학습코드에 shared cache 변수 추가

from multiprocessing import Manager

manager = Manager()
img_cache = manager.dict()
train_dataset = SkinDataset(train_data, img_cache, imsize, transform)

결과

• 현재 SSD가 아닌, HDD를 사용중이라 병목현상이 더 크게 발생했음
• batch_size 512에서 한 Epoch에 580s가 걸렸었는데, 이미지 캐싱코드를 사용하여 한 Epoch에 130s로 줄임
• I/O병목현상 해소로 GPU 사용률이 5배 이상 높아짐

Reference

• https://github.com/ptrblck/pytorch_misc/blob/master/shared_dict.py