實作訓練步驟
在這個練習中，你將在一個為影像分類任務設計的 PyTorch Lightning 模組中實作 training_step() 方法。
你的實作需要解包一個批次的影像與標籤，透過 forward 傳遞計算模型的預測，計算交叉熵損失，並記錄訓練損失。
本練習屬於課程
Scalable AI Models with PyTorch Lightning
練習說明
- 確保使用 forward 傳遞來計算預測。
- 計算交叉熵損失。
- 記錄訓練損失。
動手互動練習
試著完成這個範例程式碼，體驗一下這個練習。
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