[PyTorch] 04. validation phase를 추가해 train_model() 함수 작성하기

# 모델 생성 및 설정
net = MyModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters())

# 데이터 로드
train_data = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor())
test_data = torchvision.datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor())

# dataloader에 적재
train_loader = torch.utils.data.DataLoader(train_data, shuffle=True, batch_size=50)
test_loader = torch.utils.data.DataLoader(test_data)

# 학습(training) 진행
st = time.time()
print(f"training with {len(train_data)} data... ", end='')
for epoch in range(1):
    for x,y in train_loader:
        y_pred = net(x)
        loss = criterion(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
print(f"done (elapsed {time.time() - st}s)")

# 검증(test) 진행
with torch.no_grad():
    acc, tot = 0, 0
    for x, y in test_loader:
        y_pred = net(x)
        acc += (y==y_pred.argmax(1)).sum()
        tot += len(y)
    print(f"test accuracy : {acc}/{tot} ({100*acc/tot}%)")

import torch.utils.data as data
import torchvision
from torchvision import transforms

# 본 포스팅에서는 RGB채널이 모두 있는 CIFAR10 데이터셋을 사용한다
train_data = torchvision.datasets.CIFAR10(
  root='./data', download=True, train=True, transform=transforms.ToTensor()
)

train_size = int(len(train_data)*0.8) # 전체의 80%를 학습용으로 사용
val_size = len(train_data) - train_size
train_dataset, val_dataset = data.dataset.random_split(train_data, [train_size, val_size])

dataloaders_dict = {
  'train' : data.DataLoader(train_dataset, shuffle=True, batch_size=50),
  'val' : data.DataLoader(val_dataset, batch_size=50)
}

def train_model(net, criterion, optimizer, dataloaders_dict, num_epochs):
    # 사용 가능한 device 인식(GPU가 있으면 GPU 사용)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    net.to(device)
    print(f"Using device: {device}")
    
    # train & val 전체 진행루틴
    for epoch in range(num_epochs):
        print("-----------------")
        print(f"> Epoch {epoch+1}/{num_epochs}")
        # 각 epoch마다 'train', 'val'을 모두 진행
        for phase in ['train', 'val']:
            # 현재 phase에 따라 모델의 상태 변경
            torch.set_grad_enabled(phase == 'train')
            if phase == 'train':
                net.train()
            else:
                net.eval()
            # 첫 epoch에서는 validation만 수행(초기 상태의 성능 측정 목적)
            if (epoch == 0) and (phase == 'train'):
                continue
            # 매 epoch마다 loss(손실), correct(맞은 개수)를 기록
            epoch_loss = 0.0
            epoch_corrects = 0
            # tqdm : 진행바(progress bar)를 표시하기 위한 라이브러리
            for x,y in tqdm(dataloaders_dict[phase]):
                x, y = x.to(device), y.to(device)
                output = net(x)
                loss = criterion(output, y)
                # train phase일 때만 loss값을 기반으로 파라미터 갱신
                if phase == 'train':
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()
                # loss, correct 계산
                _, y_pred = torch.max(output, 1)
                epoch_loss += loss.item() * x.size(0) # size(0): 해당 batch의 size
                epoch_corrects += torch.sum(y_pred == y.data)
            # epoch이 끝난 후 해당 epoch에서의 평균 손실, 정확도 계산 및 출력
            epoch_loss = epoch_loss / len(dataloaders_dict[phase].dataset)
            epoch_acc = epoch_corrects.double() / len(dataloaders_dict[phase].dataset)
            print(f"{phase} Loss: {epoch_loss:.4f}  Acc: {epoch_acc:.4f} ({epoch_corrects}/{len(dataloaders_dict[phase].dataset)})")

import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import torchvision
from torchvision import models, transforms

"""
앞서 작성한 dataloaders_dict, train_model 코드 생략
"""

# resnet18 모델에서 마지막 output layer만 변경한다(총 10개의 label)
net = models.resnet18()
net.fc = nn.Linear(in_features=net.fc.in_features, out_features=10)

# criterion, optimizer 설정
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(params = net.parameters())

# 학습 진행
train_model(net, criterion, optimizer, dataloaders_dict, num_epochs=5)