始める無料で始める

LightningModule を紹介する

いよいよ最初の LightningModule を作成します。この実践演習では、分類ワークフローの基本構造を組み立てます。線形層を定義し、forward メソッドでデータをその層に通し、training step で損失を計算します。このシンプルな構造があれば、モデルの実験を始めるための確かな土台ができます。

torch と lightning.pytorch（pl としてインポート済み）はすでに読み込まれています。

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

PyTorch Lightning によるスケーラブルな AI モデルの構築

コースを見る

演習の手順

  • pl.LightningModule を継承するクラス LightModel を定義します。
  • 入力特徴量が16、出力クラスが10であることを前提に、入力を変換する線形層を定義します。

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

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

# Define the model class
class LightModel(____):
  	# Define a linear layer to transform your input
    def __init__(self):
        super().__init__()
        self.layer = ____
    def forward(self, x):
        return self.layer(x)
    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = torch.nn.functional.cross_entropy(logits, y)
        return loss
コードを編集して実行