完善 forward 方法
在 __init__ 方法中設定好各層之後,forward 方法會決定資料如何在這些層之間流動。在 PyTorch Lightning 中,這種職責分離能讓你的程式碼更乾淨、也更容易維護。你已經看過如何撰寫建構子——現在該把重點放在 forward 傳遞,確保分類邏輯清楚,並且對訓練最佳化。這裡,__init__ 中的各層已為你定義好,所以你可以專注在 forward 的流程本身。
lightning.pytorch 與 torch.nn 已分別以 pl 和 nn 匯入。
本練習屬於課程
Scalable AI Models with PyTorch Lightning
練習說明
- 在
ClassifierModel中實作forward方法。 - 在隱藏層之後套用 ReLU 啟用函式。
動手互動練習
試著完成這個範例程式碼,體驗一下這個練習。
class ClassifierModel(pl.LightningModule):
def __init__(self, input_dim, hidden_dim, output_dim):
super().__init__()
self.hidden = nn.Linear(input_dim, hidden_dim)
self.output = nn.Linear(hidden_dim, output_dim)
# Define forward method
def ____(self, ____):
# Complete the forward pass
x = self.hidden(x)
x = ____(x)
x = self.output(x)
return x