[Pytorch] 커스텀 데이터셋 만들기

커스텀 데이터셋을 만들어서 사용

딥러닝은 기본적으로 대량의 데이터를 이용하여 모델을 학습시킨다.

하지만 데이터를 한 번에 메모리에 불러와서 훈련시키면 시간과 비용 측면에서 효율적이지 않다.

따라서 데이터를 한 번에 다 부르지 않고 조금씩 나누어 불러서 사용하는 방식이 커스텀 데이터셋(custom dataset)이다.

 

CustomDataset 클래스를 구현하기 위해서는 다음과 같은 형태를 취해야 한다.

class CustomDataset(torch.utils.data.Dataset):
    def __init__(self):  # 필요한 변수를 선언하고, 데이터셋의 전처리를 해주는 함수
    
    def __len__(self):  # 데이터셋의 길이, 총 샘플의 수를 가져오는 함수
        
    def __getitem__(self, index):  # 데이터셋에서 특정 데이터를 가져오는 함수(index번째 데이터를 반환하는 함수이며, 이때 반환되는 값은 텐서의 형태를 취해야 한다.)

예시)

class CustomDataset(Dataset):
    def __init__(self, csv_file):  # cvs_file 파라미터를 통해 데이터셋을 불러온다.
        self.label = pd.read_csv(csv_file)

    def __len__(self):  # 전체 데이터셋의 크기(size)를 반환한다.
        return len(self.label)

    def __getitem__(self, idx):  # 전체 x와 y데이터 중에 해당 idx번째의 데이터를 가져온다.
        sample = torch.tensor(self.label.iloc[idx, 0:3]).int()
        label = torch.tensor(self.label.iloc[idx, 3]).int()
        return sample, label

tensor_dataset = CustomDataset('./covtype.csv')  # 데이터셋으로 covtype.csv를 사용
dataset = DataLoader(tensor_dataset, batch_size=4, shuffle=True)  # 데이터셋을 torch.utils.data.DataLoader에 파라미터로 전달한다.

torch.utils.data.DataLoader

데이터로더(DataLoader) 객체는 학습에 사용될 데이터 전체를 보관했다가 모델 학습을 할 때 배치 크기만큼 데이터를 꺼내서 사용. 이때 주의할 것은 데이터를 미리 잘라 놓는 것이 아니라 내부적으로 반복자(iterator)에 포함된 인덱스(index)를 이용하여 배치 크기만큼 데이터를 반환한다는 것.

 

데이터로더는 다음과 같이 for 문을 이용하여 구문을 반복 실행하는 것과 같다.

for i, data in enumerate(dataset, 0):
    print(i, end=' ')
    batch=data[0]
    print(batch.size())