GPUTraining & View
[yongggg's] Big-size image dataset load Tip
Yonggg
2023. 2. 1. 14:44
폴더 안에 이미지가 있을 때, 모든 이미지를 로드하면 gpu 및 ram이 터진다.
아래와 같이 dataset을 구성하면 좋다.
class MyDataset(torch.utils.data.Dataset):
def __init__(self, file_x,: str file_y: str) -> None:
super().__init__()
# file won't be loaded with mmap_mode != None
self.file_x = np.load(file_x, mmap_mode='r')
self.file_y = np.load(file_y, mmap_mode='r')
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
return (torch.from_numpy(self.file_x[idx]), torch.from_numpy(self.file_y[idx]))
def __len__(self) -> int:
return len(self.file_x)