pytorch_lightning 打印每轮的学习率

原创

mob649e816209c2 2024-09-24 07:02:59 ©著作权

©著作权归作者所有：来自51CTO博客作者mob649e816209c2的原创作品，请联系作者获取转载授权，否则将追究法律责任

如何在 PyTorch Lightning 中打印每轮学习率

引言

PyTorch Lightning 是一个高层次的深度学习框架，它简化了 PyTorch 的使用，使得模型的训练和验证过程更加规范和易于操作。了解学习率的变化对于训练过程的控制和优化至关重要。在本文中，我们将学习如何在每一轮训练中打印学习率。

流程概述

在实现打印每轮学习率的功能之前，我们先了解一下大致流程：

步骤	描述
1	创建 PyTorch Lightning 模型
2	在模型中定义学习率调度器
3	使用回调函数打印学习率
4	训练模型并观察学习率

接下来，我们通过一个以 Python 代码为主的示例来说明每一步。

流程图

下面是整个流程的可视化图示：

flowchart TD
    A[创建 PyTorch Lightning 模型] --> B[在模型中定义学习率调度器]
    B --> C[使用回调函数打印学习率]
    C --> D[训练模型并观察学习率]

步骤详细解析

1. 创建 PyTorch Lightning 模型

首先，我们需要定义一个继承自 pl.LightningModule 的模型类。

import pytorch_lightning as pl
import torch
from torch import nn

class MyModel(pl.LightningModule):
    def __init__(self, learning_rate=0.001):
        super(MyModel, self).__init__()
        # 定义网络结构
        self.layer = nn.Linear(28 * 28, 10)  # 例如：简单的全连接层
        self.learning_rate = learning_rate

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

    def training_step(self, batch, batch_idx):
        # 定义训练步骤
        x, y = batch
        y_hat = self(x)
        loss = nn.functional.cross_entropy(y_hat, y)
        return loss

    def configure_optimizers(self):
        # 配置优化器和学习率调度器
        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.1)
        return [optimizer], [scheduler]

这里定义了一个简单的网络，包含一个线性层及其损失函数。

2. 在模型中定义学习率调度器

我们在 configure_optimizers 方法中定义了优化器及其学习率调度器。在返回的列表中包含优化器和调度器，调度器会依据设定的参数自动调整学习率。

3. 使用回调函数打印学习率

为了在每一轮结束时打印当前的学习率，我们需要实现一个自定义的回调函数。

from pytorch_lightning.callbacks import Callback

class LearningRateLogger(Callback):
    def on_epoch_end(self, trainer, pl_module):
        # 获取学习率
        lr = trainer.optimizers[0].param_groups[0]['lr']
        print(f'当前学习率: {lr}')

on_epoch_end 方法被调用时，我们可以从优化器的参数组中获取并打印当前的学习率。

4. 训练模型并观察学习率

最后，我们需要将模型与回调结合起来，并开始训练。

from pytorch_lightning import Trainer

model = MyModel(learning_rate=0.01)
lr_logger = LearningRateLogger()

trainer = Trainer(callbacks=[lr_logger], max_epochs=10)
# 假设存在一个 DataLoader
# trainer.fit(model, train_dataloader=train_loader)

在初始化 Trainer 时，我们将 LearningRateLogger 作为回调传入。执行 trainer.fit(model) 时，训练过程将自动打印每轮的学习率。

甘特图

接下来是项目的时间安排，可以使用甘特图来实际展示：

gantt
    title PyTorch Lightning 学习率打印项目
    dateFormat  YYYY-MM-DD
    section 步骤
    创建模型         :a1, 2023-10-01, 2d
    定义学习率调度器 :after a1  , 2d
    定义回调函数    :after a1  , 2d
    训练模型         :after a1  , 5d