미니 배치(mini batch) [pytorch]

 

 

1. 서론

딥러닝에서 사용하는 학습 방법 중에 가장 많이 사용하는 미니 배치 사용법을 알아보도록 하겠다.

그림 1

그림 1과 같이 6개의 데이터가 존재할 때 batch size를 2라고 하면 2개의 데이터의 에러를 구하고 평균을 내어 가중치를 갱신하며 이 방식은 데이터를 매우 많이 사용하는 딥러닝에서 자주 사용되기 때문에 필수이며 pytorch에서 mini batch를 사용하는 방법을 알아보도록 하겠다.

2. 본론

mini batch를 사용하기 위해서 텐서 데이터셋과 데이터 로더가 필요하다.

from torch.utils.data import TensorDataset  # 텐서데이터셋
from torch.utils.data import DataLoader  # 데이터로더

이번 문제는 입력값의 출력을 제곱 값을 출력하는 모델을 구현하였다. 입력과 정답 쌍을 TensorDataset으로 묶은 후 DataLoader를 호출하여 데이터를 어떻게 읽어올지 선언한다. batch_size는 2로 선언하였고 데이터 순서를 섞는 shuffle을 True로 선언하였다.

# 입력
X = torch.Tensor([[1], [2], [3], [4], [5], [6], [7], [8]]).to(device)

# 정답
Y = torch.Tensor([[1], [4], [9], [16], [25], [36], [49], [64]]).to(device)

dataset = TensorDataset(X, Y)
dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

전체 코드는 아래와 같다. DataLoader를 사용하여 batch_size를 사용하면 2차원 데이터로 구성되어있던 데이터셋이 3차원으로 변형되면서 batch_size의 개수만큼 인수를 받아서 한 개의 차원을 구성한다. (8, 2) -> (4, 2, 2) 마지막 인자는 (입력, 정답)이어서 2가 된다.

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

from torch.utils.data import TensorDataset  # 텐서데이터셋
from torch.utils.data import DataLoader  # 데이터로더

device = "cuda" if torch.cuda.is_available() else "cpu"  # 장치 선택

# 입력
X = torch.Tensor([[1], [2], [3], [4], [5], [6], [7], [8]]).to(device)

# 정답
Y = torch.Tensor([[1], [4], [9], [16], [25], [36], [49], [64]]).to(device)

dataset = TensorDataset(X, Y)
dataloader = DataLoader(dataset, batch_size=2, shuffle=True)


# XOR 모델
class XOR(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = torch.nn.Linear(1, 10)
        self.layer2 = torch.nn.Linear(10, 10)
        self.layer3 = torch.nn.Linear(10, 1)

    def forward(self, x):
        x = torch.relu(self.layer1(x))
        x = torch.relu(self.layer2(x))
        x = self.layer3(x)
        return x


model = XOR().to(device)

optimizer = torch.optim.SGD(model.parameters(), lr=1e-5)

nb_epochs = 10000
for epoch in range(nb_epochs + 1):
    for batch_idx, samples in enumerate(dataloader):
        # print(batch_idx)
        # print(samples)
        x_train, y_train = samples
        # H(x) 계산
        prediction = model(x_train)

        # cost 계산
        cost = F.mse_loss(prediction, y_train)

        # cost로 H(x) 계산
        optimizer.zero_grad()
        cost.backward()
        optimizer.step()
        print(f'Epoch {epoch:4d}/{nb_epochs} Batch {batch_idx + 1}/{len(dataloader)} Cost: {cost.item():.6f}')
print(f'결과 출력: {model(X)}')

아래 그림과 같이 미니 배치 크기만큼 데이터를 넣은 후 가중치 갱신이 이루어지고 있는 것을 확인할 수 있으며 정답의 출력이 매우 정확하지는 않으나 학습이 되고 있는 것을 알 수 있었다.

3. 후기

optimizer(Adam)와 momentom, weight_decay 등 변형을 시켜서 해봤는데 이 문제는 제대로 해결 못하고 학습시킨 입력만 문제를 푸는 척(?) 하는 느낌이 든다. 선형 회귀는 나중에 더 공부를 해봐야겠다