PyTorch DataLoader Example

sklearn의 붓꽃 데이터를 활용하여 pytorch와 dataloader를 활용하여 분류 문제를 풀어 보겠습니다.

iris 데이터셋을 받아서 pandas로 데이터를 변환합니다. 변환 과정이 반드시 필요한 것은 아니지만 데이터셋을 변경하거나 학습용 컬럼 정보를 수정할 때에 도움이 됩니다.

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from sklearn.datasets import load_iris
iris = load_iris()

df = pd.DataFrame(iris.data)
df.columns = iris.feature_names
df['class'] = iris.target

다음으로 PyTorch로 데이터를 import하여 학습용 데이터를 생성합니다. 학습용 데이터는 train_data와 valid_data로 분리하되 8:2 비율로 분리합니다.

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

## Prepare Dataset
data = torch.from_numpy(df.values).float()
#data.shape = torch.Size([150, 5])

# 데이터셋에서 feature 정보와 label 데이터를 분리하여 x,y 데이터를 생성
x = data[:,:4]
y = data[:,[-1]]

# train, valid 데이터셋 분리, 데이터는 8:2 or 7:3 생성
ratio = [.8, .2]

train_cnt = int(data.size(0) * ratio[0])
valid_cnt = data.size(0) - train_cnt
print(train_cnt, valid_cnt) #120, 30

# torch.randperm을 사용해서 랜덤한 int 순열을 생성, train/valid 데이터로 분리
indices = torch.randperm(data.size(0))
x = torch.index_select(x, dim=0, index=indice).split([train_cnt, valid_cnt], dim=0)
y = torch.index_select(y, dim=0, index=indice).split([train_cnt, valid_cnt], dim=0)

pytorch에서 제공하는 Dataset과 DataLoader를 import합니다.

Dataset 클래스를 상속하여 IrisDataset 클래스를 생성하고 data, label을 입력합니다.
IrisDataset을 DataLoader에 입력하여 데이터를 batch_size 만큼 데이터를 분리하여 train_loader에 넣어줍니다.

iris 데이터셋은 총 150개 데이터입니다. 이것을 train/valid 형태로 8:2로 분리했기 때문에 train 120, valid 30개의 데이터로 각각 생성됐습니다. 이렇게 생성된 데이터를 한번에 훈련하지 않고 일정 갯수로 데이터를 묶어 줍니다. 사실 소규모의 데이터 셋에서는 이러한 batch 작업이 불필요합니다. 그러나 많은 수의 데이터를 훈련하기 위해서는 이러한 작업이 필수입니다. 이번 예제에서는 30개 단위로 묶음을 만들어보겠습니다.

파이토치에서는 이러한 묶음 작업을 할 수 있는 DataLoader라는 편리한 패키지를 제공합니다. 이러한 과정을 통해서 120개의 데이터가 30개식 4묶음으로 train_loader에 저장되게 됩니다.

from torch.utils.data import Dataset, DataLoader

# Dataset 상속
class IrisDataset(Dataset):
    
    def __init__(self, data, labels):
        super().__init__()
        self.data = data
        self.labels = labels

    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

# DataLoader
train_loader = DataLoader(dataset=IrisDataset(x[0],y[0]), batch_size=config['batch_size'], shuffle=True)
valid_loader = DataLoader(dataset=IrisDataset(x[1],y[1]), batch_size=config['batch_size'], shuffle=False)

참고로 data, train, validate, bacth_size, epoch을 이해하기 위해 예를 들어보면…
선생님이 학생들의 학력 수준을 알아보기 위해서 100문제를 만들었습니다. 선생님은 학생들에게 100문제 중에서 80 문제를 풀어보면서 수학적 원리를 설명합니다. 그러나 한번에 80문제를 풀기 어려우니 20문제씩 1~4교시 동안 풀어보게 합니다. 한번만 문제를 풀어보는 것보다는 같은 문제를 반복해서 풀어보는 것이 효과적이기 때문에 5~8교시 다시 문제를 풀어봅니다.

이제 학생들은 80문제를 20문제씩 나눠서 2번에 걸쳐 풀어본것이 됩니다. 만약 시간적 여유가 있다면 2번이 아니라 3번, 4번 풀어본다면 아마도 더 학습이 잘되겠죠.

이제 학생들이 수학원리를 잘 이해했는지 테스트해보기 위해서 남겨둔 20문제를 풀어보게 합니다. 그리고 20개의 문제를 얼마나 많은 학생이 맞췄는지를 계산해봅니다.

이러한 과정은 보통의 학습에서 매우 일반적인 방법입니다. 이제 생각해보면 100문제가 data, 80문제가 train_data, 20문제가 valid_data, 80문제를 20문제씩 나눠서 4묶음을 만드는 과정 batch, 같은 문제를 총 2회 풀어봄 epoch 이것이 지금까지의 과정에서 사용했던 용어를 정의한 것입니다.

즉, train_loader는 120개의 데이터가 30개씩 4묶음으로 되어 있는 것이 됩니다. valid_loader는 30개의 데이터가 30개씩 1묶음이 되겠네요.

자, 이제 모델을 간단히 구성합니다. 학습을 위한 모델이라기 보다는 간단히 테스트하기 위한 것임으로 간단한 모델을 만들어보겠습니다.

예측 데이터는 붓꽃의 꽃받침의 길이와 너비, 꽃잎의 길이와 너비에 따라 3종류 중 하나로 예측하는 것임으로 최종 아웃풋의 형태는 3입니다. 그리고 해당 데이터를 확률 값으로 나타내기 위하여 softmax_classification을 활용합니다.

# model 생성
model = nn.Sequential(
    nn.Linear(4,3)
)

optimizer = optim.Adam(model.parameters())

from copy import deepcopy
lowest_loss = np.inf
best_model = None
lowest_epoch = np.inf

copy 패키지로부터 deepcopy를 import합니다. 이것은 이번에 데이터를 만드는 과정과 직접적인 관련이 없기 때문에 간단히만 설명하면 객체의 모든 내용을 복사해서 새로운 하나의 객체를 만드는 것을 deep copy라고 합니다. 반대의 개념은 shallow copy 입니다.

이제 학습을 시작합니다. 이 모델은 2개의 for loop으로 되어 있습니다. 가장 먼저 나오는 for loop은 epoch에 대한 정의로 train data를 총 몇번 학습하는가에 대한 정의입니다. 다음에 나오는 또 하나의 for loop은 학습 데이터를 몇개로 나눠서 학습할 것인가 즉, batch에 대한 문제입니다.

1번 학습이 끝나면 학습의 loss를 계산해봅니다. loss는 정답과의 차이를 의미하는 것으로 작으면 작을 수록 학습이 잘됐다는 의미입니다. 한번 학습이 끝나면 valid data를 실행해봅니다. 그리고 valid에서 나온 loss와 train에서 나온 loss를 비교해보고 valid의 loss가 더 좋을 때에 해당 학습에 사용한 모델을 deepcopy해서 저장합니다.

그 이유는 무조건 학습을 오래 한다고 해서 좋은 결과가 나오는 것이 아니고 어느 순간에 학습이 정체되거나 과적합 되는 일이 있기 때문에 가장 좋은 모델을 저장하는 것입니다.

train_history, valid_history = [], []

for i in range(config['n_epochs']+1):
    model.train()
    
    train_loss, valid_loss = 0, 0
    y_hat = []
    
    # train_batch start
    for x_i, y_i in train_loader:
        y_hat_i = model(x_i)
        loss = F.cross_entropy(y_hat_i, y_i.long().squeeze())
        
        optimizer.zero_grad()
        loss.backward()

        optimizer.step()        
        train_loss += float(loss) # This is very important to prevent memory leak.

    train_loss = train_loss / len(train_loader)
    
    model.eval()
    with torch.no_grad():
        valid_loss = 0
        
        for x_i, y_i in valid_loader:
            y_hat_i = model(x_i)
            loss = F.cross_entropy(y_hat_i, y_i.long().squeeze())
            
            valid_loss += float(loss)
            
            y_hat += [y_hat_i]
            
    valid_loss = valid_loss / len(valid_loader)
    
    train_history.append(train_loss)
    valid_history.append(valid_loss)
    
    if i % config['print_interval'] == 0:
        print('Epoch %d: train loss=%.4e  valid_loss=%.4e  lowest_loss=%.4e' % (i, train_loss, valid_loss, lowest_loss))
        
    if valid_loss <= lowest_loss:
        lowest_loss = valid_loss
        lowest_epoch = i
        best_model = deepcopy(model.state_dict())
        
    model.load_state_dict(best_model)

이제 학습이 잘됐는지 아래와 같은 방법으로 train_loss와 valid_loss를 표시해봅니다.

import matplotlib.pyplot as plt

fig, loss_ax = plt.subplots()

loss_ax.plot(train_history, 'y', label='train loss')
loss_ax.plot(valid_history, 'r', label='val loss')

loss_ax.set_xlabel('epoch')
loss_ax.set_ylabel('loss')

loss_ax.legend(loc='upper left')

plt.show()

사실 이 예제는 torch의 Dataset과 DataLoader를 사용하는 방법에 대한 예제였는데 이것저것 설명하다 보니 글이 길어졌습니다.

여기서 중요한 것은 Dataset을 만들고 DataLoader를 통해서 학습에 사용하는 방법에 대한 내용이 중요하니 예제 코드를 활용해서 직접 테스트해보시기 바랍니다.