Enhancing training with Lightning callbacks

Scalable AI Models with PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

What are callbacks?

$$

  • Functions executed at key stages of training
  • Add custom actions without cluttering code
  • Enhance flexibility and control

callbacks overview

Scalable AI Models with PyTorch Lightning

What are callbacks?

$$

from lightning.pytorch.callbacks import Callback

class MyPrintingCallback(Callback):

def on_train_start(self, trainer, pl_module): print("Training is starting")
def on_train_end(self, trainer, pl_module): print("Training is ending")
  • Adding custom actions at various stages of training
Scalable AI Models with PyTorch Lightning

Lightning ModelCheckpoint callback

$$

  • Automatically saves model at specified intervals

  • Choose metric to track

  • Keep only the best model

from lightning.pytorch.callbacks
import ModelCheckpoint

checkpoint_callback = ModelCheckpoint(
monitor='val_loss',
dirpath='my/path/',
filename='{epoch}-{val_loss:.2f}',
save_top_k=1,
mode='min'
)
1 https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html
Scalable AI Models with PyTorch Lightning

Lightning EarlyStopping callback

$$

  • Monitor a metric
  • Stop training when the metric stops improving
from lightning.pytorch.callbacks
import EarlyStopping

early_stopping_callback = EarlyStopping(

monitor='val_loss',
patience=3,
mode='min'
)
1 https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.EarlyStopping.html
Scalable AI Models with PyTorch Lightning

Customizing and using lightning callbacks

from lightning.pytorch import Trainer
from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint

checkpoint = ModelCheckpoint( monitor='val_accuracy', save_top_k=2, mode='max')
early_stopping = EarlyStopping( monitor='val_accuracy', patience=5, mode='max')
trainer = Trainer(max_epochs=50, callbacks=[checkpoint, early_stopping])
1 https://lightning.ai/docs/pytorch/stable/common/trainer.html
Scalable AI Models with PyTorch Lightning

Let's practice!

Scalable AI Models with PyTorch Lightning

Preparing Video For Download...