PyTorch 모델 학습 및 평가 코드 이해하기

PyTorch를 사용하여 모델을 학습하고 평가하는 과정에서 중요한 개념과 코드를 정리한다.
특히, train()과 evaluate() 함수에서의 주요 동작을 하나씩 분석하며, 학습과 평가의 차이를 명확히 이해한다.

---

1. 학습 코드 (train 함수)

import torch
import torch.nn as nn

criterion = nn.BCELoss()
optimizer = nn.Adam(model.parameters(), lr=0.01)

def train(model, train_loader, optimizer, criterion, device):
    model.train()  # 모델을 학습 모드로 설정
    running_loss = 0.0  

    for data, label in train_loader:  
        data, label = data.to(device, non_blocking=True), label.to(device, non_blocking=True)  # GPU로 데이터 이동
        outputs = model(data)  
        loss = criterion(outputs, label)  

        optimizer.zero_grad()  # 이전 스텝의 Gradient 초기화
        loss.backward()  # Gradient 계산
        optimizer.step()  # 가중치 업데이트

        running_loss += loss.item()  # 현재 배치의 손실 값 누적

    epoch_loss = running_loss / len(train_loader)  # 평균 Loss 계산
    print(f"Epoch Loss: {epoch_loss:.4f}")
    return epoch_loss

1.1 train_loader와 미니배치 학습

train_loader는 데이터셋을 미니배치 단위로 제공하며, batch_size 설정에 따라 학습 방식이 달라진다.

• batch_size = 1 → Stochastic Gradient Descent (SGD)
• batch_size > 1 → Mini-batch Gradient Descent (일반적인 학습 방식)
• batch_size = 전체 데이터 크기 → Batch Gradient Descent

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

위 설정에서는 batch_size=32로 미니배치 학습을 진행한다.

---

1.2 model.train()의 역할

model.train()

• 모델을 학습 모드로 설정하여 Dropout과 BatchNorm이 활성화됨
• Dropout은 일부 뉴런을 랜덤하게 비활성화하여 Regularization 효과를 제공
• BatchNorm은 현재 배치의 평균과 분산을 사용하여 정규화 진행

---

1.3 역전파(Backpropagation) 과정

optimizer.zero_grad()  # 기존 Gradient 초기화
loss.backward()  # Gradient 계산
optimizer.step()  # 가중치 업데이트

1. optimizer.zero_grad() → 기존 Gradient를 초기화하여 새로운 Gradient 계산을 위한 준비
2. loss.backward() → Autograd를 사용하여 Gradient 계산
3. optimizer.step() → 계산된 Gradient를 기반으로 가중치를 업데이트

이 과정이 미니배치마다 반복되면서 모델이 학습된다.

---

2. 평가 코드 (evaluate 함수)

def evaluate(model, test_loader, criterion, device):
    model.eval()  # 평가 모드로 설정
    running_loss = 0.0  

    with torch.no_grad():  # Gradient 계산 비활성화 (메모리 절약 & 속도 향상)
        for data, label in test_loader:
            data, label = data.to(device, non_blocking=True), label.to(device, non_blocking=True)  # GPU 이동
            outputs = model(data)
            loss = criterion(outputs, label)
            running_loss += loss.item()

    epoch_loss = running_loss / len(test_loader)  # 평균 Loss 계산
    print(f"Test Loss: {epoch_loss:.4f}")
    return epoch_loss

---

2.1 model.eval()의 역할

model.eval()

• 모델을 평가 모드로 설정하여 Dropout이 비활성화됨
• BatchNorm은 학습 중 저장된 평균과 분산을 사용하여 정규화 진행

평가 모드에서는 train()과 달리 뉴런이 랜덤하게 비활성화되지 않고, 모든 뉴런이 활성화된 상태에서 예측을 수행한다.

---

2.2 torch.no_grad()의 역할

with torch.no_grad():

• Autograd 기능을 비활성화하여 Gradient 계산을 하지 않도록 설정
• 이를 통해 메모리 사용량을 줄이고, 평가 속도를 향상시킴

평가 과정에서는 가중치를 업데이트하지 않으므로, Gradient를 계산할 필요가 없다. 따라서 torch.no_grad()를 사용하면 연산량이 줄어들어 성능이 향상된다.

---

3. train() vs evaluate() 차이점

기능 train() (학습) evaluate() (평가)

model.train()	활성화 (Dropout, BatchNorm 동작)	비활성화 (Dropout, BatchNorm 고정)
model.eval()	비활성화	활성화 (Dropout 비활성, BatchNorm 고정)
torch.no_grad()	사용 안 함 (Gradient 필요)	사용 (Gradient 필요 없음)
loss.backward()	사용 (Gradient 계산)	사용 안 함
optimizer.step()	사용 (가중치 업데이트)	사용 안 함

---

4. 최종 정리

• train()에서는 model.train()을 호출하여 학습 모드로 설정하고, Gradient를 계산한 후 가중치를 업데이트
• evaluate()에서는 model.eval()을 호출하여 평가 모드로 설정하고, torch.no_grad()를 사용하여 불필요한 연산을 제거하여 성능 최적화
• train_loader의 batch_size 설정에 따라 학습 방식이 달라지며, 일반적으로 Mini-batch Gradient Descent를 사용
• running_loss를 사용하여 한 Epoch 동안의 평균 Loss를 계산하여 학습 과정을 모니터링

이제 학습과 평가 과정을 명확히 이해하고, 실제 모델을 훈련하고 검증할 수 있는 준비가 되었다.