PyTorch Lightning Callback介绍

PyTorch Lightning Callback 介绍

在 PyTorch 中，callbacks（回调函数）不是原生支持的核心功能，但在深度学习中非常常见，尤其是用来监控训练过程、调整超参数或执行特定的任务。许多高级深度学习框架（如 PyTorch Lightning 和 FastAI）都基于 PyTorch，并内置了 callback 支持。

PyTorch Lightning 提供了一个易于扩展的回调机制，允许用户在训练过程中插入自定义逻辑。回调类继承自 pytorch_lightning.callbacks.Callback，可以覆盖以下方法：

常用方法

on_fit_start: 在训练（fit）开始时调用。
on_fit_end: 在训练（fit）结束时调用。
on_train_epoch_start: 在每个训练 epoch 开始时调用。
on_train_epoch_end: 在每个训练 epoch 结束时调用。
on_validation_epoch_start: 在每个验证 epoch 开始时调用。
on_validation_epoch_end: 在每个验证 epoch 结束时调用。
on_test_epoch_start: 在测试 epoch 开始时调用。
on_test_epoch_end: 在测试 epoch 结束时调用。
on_train_batch_end: 在每个训练 batch 结束时调用。
on_validation_batch_end: 在每个验证 batch 结束时调用。
on_test_batch_end: 在每个测试 batch 结束时调用。

示例: 自定义 Callback

以下示例实现了一个打印日志的回调：

from pytorch_lightning.callbacks import Callbackclass PrintCallback(Callback):def on_train_epoch_end(self, trainer, pl_module):print(f"Epoch {trainer.current_epoch}: Training ended!")def on_validation_epoch_end(self, trainer, pl_module):print(f"Epoch {trainer.current_epoch}: Validation ended!")

使用时将回调传递给 Trainer：

from pytorch_lightning import Trainertrainer = Trainer(callbacks=[PrintCallback()])

基于 Hydra 配置实例化 Callback

Hydra 是一个灵活的配置管理工具，常用于深度学习项目中动态管理超参数。通过结合 Hydra 和 PyTorch Lightning，可以动态配置并实例化 Callback。

步骤：

1. 安装 Hydra：

pip install hydra-core --upgrade

2. 定义 Hydra 配置文件：创建一个 YAML 配置文件（如 config.yaml）来管理 Callback 的配置：

callbacks:model_checkpoint:_target_: pytorch_lightning.callbacks.ModelCheckpointmonitor: "val_loss"save_top_k: 1mode: "min"early_stopping:_target_: pytorch_lightning.callbacks.EarlyStoppingmonitor: "val_loss"patience: 5mode: "min"

3. 在代码中动态实例化：使用 hydra.utils.instantiate 方法实例化回调对象：

import hydra
from hydra.utils import instantiate
from pytorch_lightning import Trainer
from omegaconf import OmegaConf@hydra.main(config_path=".", config_name="config")
def main(cfg):# Instantiate callbacks from configcallbacks = [instantiate(cfg.callbacks[key]) for key in cfg.callbacks]# Example: Define a simple PyTorch Lightning modelfrom pytorch_lightning import LightningModuleimport torch.nn.functional as Fclass SimpleModel(LightningModule):def __init__(self):super().__init__()self.layer = torch.nn.Linear(10, 1)def forward(self, x):return self.layer(x)def training_step(self, batch, batch_idx):x, y = batchy_hat = self(x)loss = F.mse_loss(y_hat, y)return lossdef configure_optimizers(self):return torch.optim.Adam(self.parameters(), lr=0.001)# Instantiate trainertrainer = Trainer(callbacks=callbacks, max_epochs=10)# Simulated data loaderfrom torch.utils.data import DataLoader, TensorDatasetimport torchx = torch.rand(100, 10)y = torch.rand(100, 1)train_loader = DataLoader(TensorDataset(x, y), batch_size=32)model = SimpleModel()trainer.fit(model, train_loader)if __name__ == "__main__":main()