Pytorch-Lightning 是 pytorch 的抽象和包装。它进一步简化了 pytorch 中的训练代码，将一个网络的训练、验证抽象为一个更简单的类。

基本类型

LightningModule

pytorch lightning 的基类，继承自 pytorch 的 nn.Module，用于构建其他 pytorch 网络模型。

由于继承自 nn.Module ，这个网络也可以指定 forward 函数用于前向计算。训练过程主要的方法是 training_step 和 configuration_optimizer，前者负责处理一步训练时如何优化函数，后者指定训练过程使用的优化器。

方法

以下方法都为虚函数，需要在继承类中自己定义。

training_step 函数接受两个参数：batch 和 batch_idx，batch 是 Dataloader 一次循环给出的一个 batch，batch_idx 则为编号。

class MyModule(pl.LightningModule):
    def __init__(self, ...):
        self.net = UNet() # 指定一个网络
        self.mse = nn.MSELoss()

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

configure_optimizers 则要求返回一个指定学习率的优化器：

class MyModule(pl.LightningModule):
    # ... 省略其他函数
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.net.parameters(), lr=1e-3)
        return optimizer

接口

文档原文称为 hook，是一个类中的一些函数接口，类似于 C++ 中的函数指针。

使用该类时，这些接口可以挂载其他函数，使该类实现更为丰富的功能。

Trainer

训练基类，用于整合网络模块和数据集，并使网络在指定数据集上训练。

最基本地，使用 fit 方法，可以拟合一个给定的模型：

trainer = Trainer()
trainer.fit(model, train_dataloader, val_dataloader)

fit 方法也可以直接传入 Dataset，此时将基于默认的或指定的聚合方法，自动产生相应的 Dataloader，将 Dataset 取出的元素随机顺序排列，并且整合为一个 batch。

参数

Trainer 的参数非常多，主要包括：

accelerator 和 devices：指定训练的加速硬件和硬件数量。个人训练或者科研通常为 "gpu" 和 1，即 trainer = Trainer(accelerator="gpu", devices=1)

gpu：直接指定 GPU 加速以及 GPU 数量，单卡训练时设置 gpu=1。gpu=None 或 gpu=0 表示不使用 GPU。

callbacks：指定一个回调函数（Callback 类的实例），或者回调函数列表。这些函数将在单步训练的之前或之后被调用。

check_val_every_n_epoch：每经过 n 个 epoch 之后才验证一次，默认为 1。

max_epochs：训练轮数，不指定时默认为 1000，即默认 1000 个 epochs。

precision：浮点精度，设置为字符串的 "32" 时代表全程使用 32 位浮点数（float），设置为 "64" 时表示全程使用 double。还支持混合精度浮点数。

fast_dev_run：设置为 True 时，只执行一轮训练和验证。一般用于 debug，检验代码有无错误。

技巧

ckpt 的保存与加载

每隔 n 个 epochs 保存一次 ckpt

直接使用 ModelCheckpoint 回调函数即可，使用 every_n_epochs 指定每隔几个 epochs 保存一次。

checkpoint_callback = ModelCheckpoint(
    dirpath=os.path.join(ckpt_save_path, 'ddpm_flow_match/'),   # 保存路径
    every_n_epochs=100,       # 每隔100 epoch保存
    save_top_k=-1,            # 保存所有符合条件的模型
    filename='model-{epoch}-loss-{loss:.3f}'  # 文件名格式
)

trainer = pl.Trainer(
    max_epochs=num_epochs,
    accelerator="gpu",  # 自动选择可用加速器（GPU/TPU等）
    devices=[1],
    callbacks=[checkpoint_callback]
)

加载保存的 ckpt

在 trainer.fit 方法中，可以直接使用 ckpt_path 参数，从过去的 ckpt 加载参数和优化器状态：

trainer.fit(
    pl_model, 
    train_dataloader,
    ckpt_path=last_ckpt_path
)