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)