lightning/examples/pytorch/basics

Basic Examples

Use these examples to test how Lightning works.

AutoEncoder

This script shows you how to implement a CNN auto-encoder.

# CPU
python autoencoder.py

# GPUs (any number)
python autoencoder.py --trainer.accelerator 'gpu' --trainer.devices 2

# Distributed Data Parallel (DDP)
python autoencoder.py --trainer.accelerator 'gpu' --trainer.devices 2 --trainer.strategy 'ddp'

---

Backbone Image Classifier

This script shows you how to implement a LightningModule as a system. A system describes a LightningModule which takes a single torch.nn.Module which makes exporting to producion simpler.

# CPU
python backbone_image_classifier.py

# GPUs (any number)
python backbone_image_classifier.py --trainer.accelerator 'gpu' --trainer.devices 2

# Distributed Data Parallel (DDP)
python backbone_image_classifier.py --trainer.accelerator 'gpu' --trainer.devices 2 --trainer.strategy 'ddp'

---

Transformers

This example contains a simple training loop for next-word prediction with a Transformer model on a subset of the WikiText2 dataset.

python transformer.py

---

PyTorch Profiler

This script shows you how to activate the PyTorch Profiler with Lightning.

python profiler_example.py