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.
This exercise is part of the course
Scalable AI Models with PyTorch Lightning
Exercise instructions
- 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.
Hands-on interactive exercise
Have a go at this exercise by completing this sample code.
# 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