トレーニングステップの実装
この演習では、画像分類タスク向けの PyTorch Lightning モジュールで training_step() メソッドを実装します。
実装では、画像とラベルのバッチを展開し、フォワードパスでモデルの予測値を計算し、クロスエントロピー損失を算出したうえで、トレーニング損失をログに記録します。
この演習はコースの一部です
PyTorch Lightning によるスケーラブルな AI モデルの構築
演習の手順
- フォワードパスを使って予測値を計算してください。
- クロスエントロピー損失を計算してください。
- トレーニング損失をログに記録してください。
実践的なインタラクティブ演習
このサンプルコードを完成させて、この演習に挑戦してみましょう。
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