Lightning コールバックによるトレーニングの強化

PyTorch Lightning によるスケーラブルな AI モデルの構築

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")

トレーニングのさまざまな段階でカスタム処理を追加

Lightning の ModelCheckpoint コールバック

指定した間隔でモデルを自動保存
追跡する指標を選択
最良のモデルのみを保持

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'

)

¹ https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html

Lightning の EarlyStopping コールバック

指標を監視
指標の改善が止まったらトレーニングを停止

from lightning.pytorch.callbacks
import EarlyStopping

early_stopping_callback = EarlyStopping(

    monitor='val_loss',

    patience=3,

    mode='min'

)

¹ https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.EarlyStopping.html

Lightning コールバックのカスタマイズと使用

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])

¹ https://lightning.ai/docs/pytorch/stable/common/trainer.html

練習しましょう！

PyTorch Lightning によるスケーラブルな AI モデルの構築