Scalable AI Models with PyTorch Lightning
Sergiy Tkachuk
Director, GenAI Productivity
$$
$$
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")
$$
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'
)
$$
from lightning.pytorch.callbacks import EarlyStopping early_stopping_callback = EarlyStopping(
monitor='val_loss',
patience=3,
mode='min'
)
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])
Scalable AI Models with PyTorch Lightning