[python] 파이썬 PyTorch에서 학습 중 모델의 가중치를 업데이트하는 방법은?

22 Nov 2023

python

옵티마이저를 사용하여 모델의 가중치를 업데이트하는 코드는 다음과 같습니다:

import torch
import torch.nn as nn
import torch.optim as optim

# 모델 생성
model = nn.Sequential(
    nn.Linear(input_dim, hidden_dim),
    nn.ReLU(),
    nn.Linear(hidden_dim, output_dim)
)

# 손실 함수 설정
criterion = nn.CrossEntropyLoss()

# 옵티마이저 설정
optimizer = optim.SGD(model.parameters(), lr=learning_rate)

# 모델 학습 반복
for epoch in range(num_epochs):
    # Forward 단계
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    
    # Backward 단계
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # 에폭(epoch)마다 손실(loss) 출력
    if (epoch+1) % print_interval == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item()}')

위의 코드에서 model.parameters()는 모델의 가중치들을 반환합니다. 이를 옵티마이저에 전달하여 가중치를 업데이트하게 됩니다. 이후 옵티마이저의 step() 함수를 호출하여 가중치를 실제로 업데이트합니다.

학습 과정에서는 Forward 단계, Backward 단계를 반복하며 손실 함수의 값을 최소화하기 위해 모델의 가중치를 업데이트합니다. 각 에폭마다 손실 값이 출력되도록 설정할 수 있습니다.

자세한 내용은 PyTorch 공식 문서를 참고하시기 바랍니다.