[DL] 파이토치 기본 학습 과정

간단한 파이토치 전체 학습 과정

라이브러리

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F # torch.nn 안의 클래스들을 사용해도 무방
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import transforms

학습 디바이스 설정

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

하이퍼 파라미터 설정

input_size = 784
num_classes = 10
learning_rate = 0.001
batch_size = 64
num_epochs = 1

데이터셋, 데이터 로더 설정

MNIST 사용

train_dataset = datasets.MNIST(root='dataset/',
                               train=True,
                               transform=transforms.ToTensor(),
                               download=True)
train_loader = DataLoader(dataset=train_dataset,
                          batch_size=batch_size,
                          shuffle=True)

test_dataset = datasets.MNIST(root='dataset/',
                               train=False,
                               transform=transforms.ToTensor(),
                               download=True)
test_loader = DataLoader(dataset=test_dataset,
                          batch_size=batch_size,
                          shuffle=True)

모델 설정

class LinearNN(nn.Module):
    def __init__(self, input_size, num_classes): # 28x28
        super().__init__()
        self.fc1 = nn.Linear(input_size, 50)
        self.fc2 = nn.Linear(50, num_classes)
    
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

모델 생성 및 손실 함수, 옵티마이저 설정

model = LinearNN(input_size=input_size, num_classes=num_classes).to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

모델 학습

for epoch in range(num_epochs):

    # Training
    train_loss, train_acc = 0, 0
    model.train()
    for batch_idx, (X_train, y_train) in enumerate(train_loader):
        # Set device
        X_train = X_train.to(device)
        y_train = y_train.to(device)

        # Set correct shape
        X_train = X_train.reshape(X_train.shape[0], -1) # [64, 1, 28, 28] -> [64, 784]

        # forward
        y_proba = model(X_train)
        loss = criterion(y_proba, y_train)

        # backward
        optimizer.zero_grad()
        loss.backward()

        # update weights
        optimizer.step()

        train_loss += loss.item()
        _, y_pred = torch.max(y_proba, axis=1)
        train_acc += sum(y_pred == y_train).item() / len(y_train)
    
    train_loss /= len(train_loader)
    train_acc /= len(train_loader)

    # Test
    test_loss, test_acc = 0, 0
    model.eval()
    with torch.inference_mode():
        for X, y in test_loader:
            X = X.to(device)
            y = y.to(device)

            X = X.reshape(X.shape[0], -1)

            test_proba = model(X)
            test_loss += criterion(test_proba, y)
            _, test_pred = torch.max(test_proba, axis=1)

            test_acc += sum(test_pred == y).item() / len(y)

        test_loss /= len(test_loader)
        test_acc /= len(test_loader)

    print(f"Epochs: {epoch+1:2} | "
          f"Train Loss: {train_loss:.5f} | "
          f"Train Acc: {train_acc:.5f} | "
          f"Test Loss: {test_loss:.5f} | "
          f"Test Acc: {test_acc:.5f} | ")