基于PyTorch Lightning的学习率打印方案

项目背景

在深度学习的训练过程中，学习率是一个极为重要的超参数，直接影响到模型的收敛速度和最终性能。PyTorch Lightning是一个高度模块化的深度学习框架，其在保持PyTorch灵活性的同时，还提供了许多便捷的功能。其中之一就是通过 Trainer 对象管理训练过程。在训练期间，监控和打印学习率能够让研究者更好地理解模型训练情况，及时调整训练策略。

方案概述

本方案旨在展示如何在使用PyTorch Lightning的 Trainer 中打印学习率，包括在每个训练周期后和训练过程中实时打印，并最终通过可视化展示学习率变化情况。

环境准备

首先，请确保你已经安装了PyTorch和PyTorch Lightning。可以通过以下命令安装：

pip install torch pytorch-lightning

示例代码结构

接下来，我们将创建一个简单的神经网络，并使用PyTorch Lightning进行训练。以下是基本的代码结构：

import torch
import torch.nn as nn
import pytorch_lightning as pl
from torch.optim import Adam

class SimpleModel(pl.LightningModule):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.layer = nn.Linear(28 * 28, 10)

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

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = nn.CrossEntropyLoss()(y_hat, y)
        return loss

    def configure_optimizers(self):
        optimizer = Adam(self.parameters(), lr=0.001)
        return optimizer

打印学习率

为了打印学习率，我们可以在 on(epoch_end) 方法中实现，另外我们还可以记录每个epoch的学习率，用于后续的可视化分析。以下是完整的代码示例：

class LitModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = SimpleModel()
        self.lr = 0.001
        self.learning_rates = []

    def training_step(self, batch, batch_idx):
        return self.model.training_step(batch, batch_idx)

    def configure_optimizers(self):
        optimizer = Adam(self.model.parameters(), lr=self.lr)
        return optimizer

    def on_epoch_end(self):
        current_lr = self.trainer.optimizers[0].param_groups[0]['lr']
        self.learning_rates.append(current_lr)
        print(f'学习率：{current_lr:.6f}')

# 实际的训练过程
trainer = pl.Trainer(max_epochs=10)
model = LitModel()

# 假设有一个训练数据加载器train_loader
train_loader = ...
trainer.fit(model, train_loader)

在上面的代码中，我们重载了 on_epoch_end 函数来打印每个epoch结束时的学习率。每次调用训练步骤后，学习率会被存储到 self.learning_rates 列表中。

可视化学习率变化

最后，我们将学习率数据可视化，以便更直观地理解其变化。以下是使用Matplotlib绘制学习率变化曲线的代码示例：

import matplotlib.pyplot as plt

# 假设训练模型已经完成
plt.figure(figsize=(10, 5))
plt.plot(model.learning_rates, marker='o')
plt.title('学习率变化情况')
plt.xlabel('Epoch')
plt.ylabel('学习率')
plt.grid()
plt.show()

饼状图展示

此外，我们在代码中可以添加一个饼状图，用于展示学习率在不同阶段下的占比情况。下面是饼状图的示例代码，使用Mermaid语法表示：

pie
    title 学习率占比
    "初始学习率": 40
    "优化后学习率": 30
    "调度后学习率": 30

结论

通过此方案，我们实现了在PyTorch Lightning框架中实时打印学习率，并对学习率变化进行了可视化。通过这种方式，我们可以更好地理解模型的训练过程，并在适当的时候调整学习率策略。希望此方案能帮助正在进行深度学习研究的同行们有效监控他们的模型训练过程。