[PyTorch] Optimizer 및 Training

Optimization & Training

https://github.com/pytorch/examples/tree/master/mnist

Load Packages

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms

import numpy as np

no_cuda = False

# cuda를 사용할지 안할지
use_cuda = not no_cuda and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")

=> device(type='cuda')

Preprocess

seed = 1
batch_size = 64
test_batch_size = 64

torch.manual_seed(seed)

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('dataset', train=True, download=True,
                transform=transforms.Compose([
                  transforms.ToTensor(),
                  transforms.Normalize((0.1307,), (0.3081,))
                ])),
                batch_size=batch_size, shuffle=True)


test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('dataset', train=False,
                transform=transforms.Compose([
                  transforms.ToTensor(),
                  transforms.Normalize((0.1307,), (0.3081,))
                ])),
                batch_size=test_batch_size, shuffle=True)

Model

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4 * 4 * 50, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = x.view(-1, 4 * 4 * 50)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

Optimization

Model과 Optimization를 설정합니다.

SGD 사용

• 확률적 경사 하강법(Stochastic Gradient Descent, SGD) 옵티마이저
• Lr(learning Rate) : 0 보다 크거나 같은 float 값. 학습률
• momentum: 0 보다 크거나 같은 float 값. SGD를 적절한 방향으로 가속화하며, 흔들림(진동)을 줄여주는 매개변수

model = Net().to(device)
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.5)

parameters를 확인합니다.

params = list(model.parameters())
for i in range(8):
    print(params[i].size())

# Weight, Bais를 순서대로 보여준다.
torch.Size([20, 1, 5, 5])
torch.Size([20])
torch.Size([50, 20, 5, 5])
torch.Size([50])
torch.Size([500, 800])
torch.Size([500])
torch.Size([10, 500])
torch.Size([10])

Before Training

학습하기 전에 Model이 Train할 수 있도록 Train Mode로 변환합니다.
Convolution 또는 Linear 뿐만 아니라, DropOut과 Batch Normalization과 같이 parameter를 가진 Layer들도 학습하기 위해 준비합니다.

# train mode
model.train()

# Out
Net(
  (conv1): Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1))
  (conv2): Conv2d(20, 50, kernel_size=(5, 5), stride=(1, 1))
  (fc1): Linear(in_features=800, out_features=500, bias=True)
  (fc2): Linear(in_features=500, out_features=10, bias=True)
)

모델에 넣기 위한 첫 Batch 데이터를 추출하고 cpu 또는 gpu device에 컴파일합니다.

data, target = next(iter(train_loader))
data, target = data.to(device), target.to(device)
data.shape, target.shape

=> (torch.Size([64, 1, 28, 28]), torch.Size([64]))

기울기(gradients)를 clear 해서 새로운 최적화 값을 찾기 위해 준비합니다. 그리고 준비한 데이터를 model 에 넣어 output 을 얻습니다. Model 에서 예측한 결과를 Loss Function 에 넣습니다.

optimizer.zero_grad()
output = model(data)

# Negative Log-Likelihood Loss 사용
loss = F.nll_loss(output, target)

Back Propagation 을 통해 기울기를 계산합니다. 계산된 기울기는 Parameter 에 업데이트를 합니다.

# 기울기 계산
loss.backward()
# Update
optimizer.step()

Start Training

위의 최적화 과정을 반복하여 학습을 시작합니다.

epochs = 1
log_interval = 100

for epoch in range(1, epochs + 1):
    # Train Mode
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()

        if batch_idx % log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100 * batch_idx / len(train_loader), loss.item()
            ))

# Out
Train Epoch: 1 [0/60000 (0%)] Loss: 2.290735
Train Epoch: 1 [6400/60000 (11%)] Loss: 2.228956
Train Epoch: 1 [12800/60000 (21%)]  Loss: 2.121080
Train Epoch: 1 [19200/60000 (32%)]  Loss: 1.893549
Train Epoch: 1 [25600/60000 (43%)]  Loss: 1.570034
Train Epoch: 1 [32000/60000 (53%)]  Loss: 1.213198
Train Epoch: 1 [38400/60000 (64%)]  Loss: 0.993993
Train Epoch: 1 [44800/60000 (75%)]  Loss: 0.778737
Train Epoch: 1 [51200/60000 (85%)]  Loss: 0.732925
Train Epoch: 1 [57600/60000 (96%)]  Loss: 0.624502