PyTorch에서는 데이터를 쉽게 다루게 하기 위해 torch.utils.data.Dataset과 torch.utils.data.DataLoader 클래스를 제공합니다. Dataset 클래스를 사용하여 사용자 데이터셋을 정의할 수 있으며 이를 통해 데이터를 불러오고 전처리하는 과정을 학습 코드로부터 분리하여 관리할 수 있습니다. DataLoader 클래스는 데이터셋을 배치 단위로 쉽게 접근할 수 있도록 iterable한 객체로 변환시켜주는 역할을 합니다.
Dataset
Dataset 클래스를 상속 받아 필요한 사용자 데이터셋을 만들어 줄 수 있습니다. Dataset 클래스의 구성은 다음과 같습니다.
- __init__: 객체가 생성될 때 데이터를 불러오고 필요한 변수를 선언할 수 있습니다.
- __len__: 데이터셋의 샘플 개수를 반환할 때 사용됩니다.
- __getitem__: 데이터셋에서 idx에 해당하는 샘플을 불러옵니다.
import torch
from torch.utils.data import Dataset
class Custom_Dataset(Dataset):
def __init__(self):
# 데이터를 불러옵니다.
self.data_x = [[5, 75, 812], [3, 61, 655]]
self.data_y = [[0], [1]]
def __len__(self):
# 데이터셋의 샘플 개수를 반환합니다.
return len(self.data_x)
def __getitem__(self, idx):
# idx에 해당하는 샘플을 반환합니다.
x = torch.Tensor(self.data_x[idx])
y = torch.Tensor(self.data_y[idx])
return x, y
train_dataset = CustomDataset()
Dataset 클래스를 이용해 다양한 처리가 가능한데 아래와 같이 train data과 test data을 각각 나눠서 생성해 줄 수도 있습니다. 또한 데이터 전처리가 필요할 경우 __init__() 또는 __getitem__() 에서 수행될 수 있습니다.
import pandas as pd
class Titanic_Dataset(Dataset):
def __init__(self, flag):
# 데이터를 불러옵니다.
df_raw = pd.read_csv('titanic.csv')
# 7:3의 비율로 train data와 test data를 나눠서 가져옵니다.
if flag == 'train':
df_data = df_raw[:int(len(df_raw) * 0.7)]
elif flag == 'test':
df_data = df_raw[int(len(df_raw) * 0.7):]
# feature로 3개 변수, label로 Survived 변수를 선택합니다.
self.data_x = df_data[['Pclass','Age','Parch']].values
self.data_y = df_data[['Survived']].values
def __len__(self):
# 데이터셋의 샘플 개수를 반환합니다.
return len(self.data_x)
def __getitem__(self, idx):
# idx에 해당하는 샘플을 반환합니다.
x = torch.Tensor(self.data_x[idx])
y = torch.Tensor(self.data_y[idx])
return x, y
train_dataset = Titanic_Dataset(flag='train')
test_dataset = Titanic_Dataset(flag='test')
DataLoader
DataLoader 클래스는 데이터셋을 학습에 용이한 mini-batch 형태로 변환시켜주는 역할을 합니다. DataLoader는 dataset을 인자로 받으며 매개변수는 아래와 같습니다.
DataLoader(dataset, batch_size=1, shuffle=False, sampler=None,
batch_sampler=None, num_workers=0, collate_fn=None,
pin_memory=False, drop_last=False, timeout=0,
worker_init_fn=None, *, prefetch_factor=2,
persistent_workers=False)
- batch_size: 배치 하나당 들어갈 샘플의 수
- shuffle: 데이터 순서를 무작위로 변경할지 여부
- drop_last: 데이터셋을 나누다보면 마지막 배치는 batch_size와 다를 수 있는데 이를 드랍할지 여부
DataLoader는 iterable한 객체를 반환하므로 for문과 결합하여 아래와 같이 사용될 수 있습니다. DataLoader를 이용하여 매 iteration 마다 반환된 batch_x, batch_y를 통해 딥러닝 모델을 학습할 수 있습니다.
from torch.utils.data import DataLoader
# 1) 앞에서 정의한 Custom_Dataset을 통해 train_dataset을 생성합니다.
train_dataset = Custom_Dataset()
# 2) DataLoader를 통해 train_dataset를 size=1 배치로 구성된 iterable 객체로 변환합니다.
train_loader = DataLoader(train_dataset, batch_size=1, shuffle=False, drop_last=False)
# 3) DataLoader는 Dataset의 __getitem__() 메서드를 통해 iteration마다 x, y를 반환합니다.
for n, (batch_x, batch_y) in enumerate(train_loader):
print('[ Iter', n, '] batch_x:', batch_x, 'batch_y:', batch_y)
실행 결과
[ Iter 0 ] batch_x: tensor([[ 5., 75., 812.]]) batch_y: tensor([[0.]])
[ Iter 1 ] batch_x: tensor([[ 3., 61., 655.]]) batch_y: tensor([[1.]])
오타나 잘못된 부분 있으면 댓글 부탁드립니다. 도움이 되셨다면 공감 눌러 주시면 감사하겠습니다 :)
댓글