認識 LightningModule
準備好建立你的第一個 LightningModule 吧！在這個動手做的練習中，你會建立分類流程的核心結構。你會定義一個線性層，在 forward 方法中將資料傳入，並在訓練步驟中計算損失。這樣清楚的結構能為你後續實驗各種模型打下穩固基礎。
torch 與 lightning.pytorch（匯入為 pl）已為你預先載入。
本練習屬於課程
Scalable AI Models with PyTorch Lightning
練習說明
- 定義一個類別
LightModel，並繼承自pl.LightningModule。 - 定義一個線性層來轉換輸入，假設輸入特徵為 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