始める無料で始める

Implementing the training step

In this exercise, you'll implement the training_step() method in a PyTorch Lightning module designed for an image classification task. Your implementation should unpack a batch of images and labels, compute the model predictions via the forward pass, calculate the cross entropy loss, and log the training loss.

この演習はコースの一部です

Scalable AI Models with PyTorch Lightning

コースを見る

演習の手順

  • Ensure that you compute predictions using the forward pass.
  • Calculate the cross entropy loss.
  • Log the training loss.

実践的なインタラクティブ演習

このサンプルコードを完成させて、この演習に挑戦してみましょう。

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
コードを編集して実行