Introducing the LightningModule
Get ready to build your first LightningModule! In this hands-on exercise, you'll set up the core structure of a classification workflow. You'll define a linear layer, pass data through it in the forward method, and compute the loss in the training step. This clean structure gives you a solid foundation to start experimenting with your models.
The torch and lightning.pytorch, imported as pl, have been preloaded for you.
本练习是课程的一部分
Scalable AI Models with PyTorch Lightning
练习说明
- Define a class
LightModelthat inherits frompl.LightningModule. - Define a linear layer to transform your input, assuming the input features are 16 and there are 10 output classes.
交互式实操练习
通过完成这段示例代码来试试这个练习。
# 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