training step 구현하기
이 연습 문제에서는 이미지 분류 작업을 위한 PyTorch Lightning 모듈에서 training_step() 메서드를 구현해 보겠습니다.
배치에서 이미지와 레이블을 언패킹하고, forward 패스를 통해 모델 예측을 계산한 다음, cross entropy 손실을 구하고, 학습 손실을 로그로 남기세요.
이 연습은 강의의 일부입니다
PyTorch Lightning으로 만드는 확장 가능한 AI 모델
연습 안내
- forward 패스를 사용해 예측을 계산하세요.
- cross entropy 손실을 계산하세요.
- 학습 손실을 로그로 기록하세요.
실습형 인터랙티브 연습
이 예제를 이 샘플 코드를 완성하여 풀어보세요.
from torch.nn.functional import cross_entropy
def training_step(self, batch, batch_idx):
x, y = batch
# Ensure that you compute predictions using the forward pass
y_hat = ____
# Calculate the cross entropy loss
loss = ____
# Log the loss
self.____("train_loss", loss)
return loss