• Docs >
  • Build and Train a Model

Build and Train a Model

Required background: Basic Python familiarity and complete the guide.

Goal: We’ll walk you through the creation of a model using PyTorch Lightning.

A simple PyTorch Lightning script

Let’s assume you already have a folder with those two files.

    train.py            # your own script to train your models
    requirements.txt    # your python requirements.

If you don’t, simply create a pl_project folder with those two files and add the following PyTorch Lightning code in the train.py file. This code trains a simple AutoEncoder on MNIST Dataset.

import os

import torch
import torch.nn.functional as F
from torch import nn
from torch.utils.data import DataLoader, random_split
from torchvision import transforms as T
from torchvision.datasets import MNIST

import pytorch_lightning as pl

class LitAutoEncoder(pl.LightningModule):
    def __init__(self):
        self.encoder = nn.Sequential(nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 3))
        self.decoder = nn.Sequential(nn.Linear(3, 128), nn.ReLU(), nn.Linear(128, 28 * 28))

    def forward(self, x):
        # in lightning,
        # forward defines the prediction/inference actions
        embedding = self.encoder(x)
        return embedding

    def training_step(self, batch, batch_idx):
        # training_step defines the train loop.
        # It is independent of forward
        x, y = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        self.log("train_loss", loss)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

dataset = MNIST(os.getcwd(), download=True, transform=T.ToTensor())
train, val = random_split(dataset, [55000, 5000])

autoencoder = LitAutoEncoder()
trainer = pl.Trainer(accelerator="auto")
trainer.fit(autoencoder, DataLoader(train), DataLoader(val))

Add the following to the requirements.txt file.


Simply run the following commands in your terminal to install the requirements and train the model.

pip install -r requirements.txt
python train.py

Get through PyTorch Lightning Introduction to learn more.

Next Steps