lightning的hook顺序

结果

setup: 训练循环开始前设置数据加载器和模型。

configure_optimizers: 设置优化器和学习率调度器。

on_fit_start: 训练过程开始。

on_train_start: 训练开始。

on_train_epoch_start: 每个训练周期开始。

on_train_batch_start: 每个训练批次开始。

on_before_backward: 反向传播之前。

on_after_backward: 反向传播之后。

on_before_zero_grad: 清空梯度之前。

on_after_zero_grad: 清空梯度之后。

on_before_optimizer_step: 优化器步骤之前。

on_train_batch_end: 每个训练批次结束。

on_train_epoch_end: 每个训练周期结束。

on_train_end: 训练结束。

on_fit_end: 训练过程结束。

测试代码

py 复制代码
import torch
from torch.utils.data import DataLoader, TensorDataset
from pytorch_lightning import LightningModule, Trainer, Callback
from pytorch_lightning.callbacks import Callback

# 定义一个简单的线性回归模型
class LinearRegression(LightningModule):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(1, 1)

    def forward(self, x):
        return self.linear(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = torch.nn.functional.mse_loss(y_hat, y)
        return loss
    
    def on_after_backward(self, *args, **kwargs):
        print("After backward is called!", args, kwargs)
        return super().on_after_backward(*args, **kwargs)
    
    def on_before_zero_grad(self, *args, **kwargs):
        print("Before zero grad is called!", args, kwargs)
        return super().on_before_zero_grad(*args, **kwargs)
    
    def on_after_zero_grad(self, *args, **kwargs):
        print("After zero grad is called!", args, kwargs)
        return super().on_after_zero_grad(*args, **kwargs)
    
    def on_before_backward(self, *args, **kwargs):
        print("Before backward is called!", args, kwargs)
        return super().on_before_backward(*args, **kwargs)
    
    def on_before_optimizer_step(self, *args, **kwargs):
        print("Before optimizer step is called!", args, kwargs)
        return super().on_before_optimizer_step(*args, **kwargs)
    
    def on_after_optimizer_step(self, *args, **kwargs):
        print("After optimizer step is called!", args, kwargs)
        return super().on_after_optimizer_step(*args, **kwargs)
    
    def on_fit_start(self, *args, **kwargs):
        print("Fit is starting!", args, kwargs)
        return super().on_fit_start(*args, **kwargs)
    
    def on_fit_end(self, *args, **kwargs):
        print("Fit is ending!", args, kwargs)
        return super().on_fit_end(*args, **kwargs)
    
    def setup(self, *args, **kwargs):
        print("Setup is called!", args, kwargs)
        return super().setup(*args, **kwargs)
    
    def configure_optimizers(self, *args, **kwargs):
        print("Configure Optimizers is called!", args, kwargs)
        return super().configure_optimizers(*args, **kwargs)
    
    def on_train_start(self, *args, **kwargs):
        print("Training is starting!", args, kwargs)
        return super().on_train_start(*args, **kwargs)
    
    def on_train_end(self, *args, **kwargs):
        print("Training is ending!", args, kwargs)
        return super().on_train_end(*args, **kwargs)
    
    def on_train_batch_start(self, *args, **kwargs):
        print(f"Training batch is starting!", args, kwargs)
        return super().on_train_batch_start(*args, **kwargs)
    
    def on_train_batch_end(self, *args, **kwargs):
        print(f"Training batch is ending!", args, kwargs)
        return super().on_train_batch_end(*args, **kwargs)
    
    def on_train_epoch_start(self, *args, **kwargs):
        print(f"Training epoch is starting!", args, kwargs)
        return super().on_train_epoch_start(*args, **kwargs)
    
    def on_train_epoch_end(self, *args, **kwargs):
        print(f"Training epoch is ending!", args, kwargs)
        return super().on_train_epoch_end(*args, **kwargs)
    
    
# 创建数据集
x_train = torch.tensor([[1.0], [2.0], [3.0], [4.0]], dtype=torch.float)
y_train = torch.tensor([[2.0], [4.0], [6.0], [8.0]], dtype=torch.float)
train_dataset = TensorDataset(x_train, y_train)
train_loader = DataLoader(train_dataset, batch_size=2)

# 创建模型和训练器
model = LinearRegression()
trainer = Trainer(max_epochs=2)

# 开始训练
trainer.fit(model, train_loader)
相关推荐
Debug_Snail20 小时前
【营销策略算法】关联规则学习-购物篮分析
大数据·人工智能
山烛20 小时前
深度学习:残差网络ResNet与迁移学习
人工智能·python·深度学习·残差网络·resnet·迁移学习
CareyWYR20 小时前
每周AI论文速递(250901-250905)
人工智能
eleqi20 小时前
Python+DRVT 从外部调用 Revit:批量创建梁(2)
python·系统集成·revit·自动化生产流水线·外部访问
BYSJMG21 小时前
计算机毕设大数据方向:基于Spark+Hadoop的餐饮外卖平台数据分析系统【源码+文档+调试】
大数据·hadoop·分布式·python·spark·django·课程设计
亚里随笔21 小时前
VERLTOOL:打通LLM工具强化学习的“任督二脉”,实现多模态多任务统一训练
人工智能·语言模型·llm·agentic
API流转日记21 小时前
对接gemini-2.5-flash-image-preview教程
人工智能·gpt·ai·chatgpt·ai作画
OC溥哥99921 小时前
Flask论坛与个人中心页面开发教程完整详细版
后端·python·flask·html
Rhys..21 小时前
python + Flask模块学习 2 接收用户请求并返回json数据
python·学习·flask
说私域21 小时前
开源AI智能名片链动2+1模式S2B2C商城小程序服务提升复购率和转介绍率的研究
人工智能·小程序