Pytorch Lightning框架笔记（一）：模型的基本使用【LightningModule】

传统的pytorch库是存在很大缺陷的，如：半精度训练、BatchNorm同步、单机多卡训练等
为了一站式解决这些问题，提出了Pytorch Lightning框架。

Pytorch Lightning流程

Pipeline

PineLine，有一个固定的顺序：

• def init(self)
• 训练training_step(self, batch, batch_idx)
• 校验validation_step(self, batch, batch_idx)
• 测试 test_step(self, batch, batch_idx)

主要是实现这四个函数的重写
还有一些其他的函数，为了方便实现其他的一些功能，更为完整的流程：

• ● training_step后加入
  training_step_end(self，batch_parts)和training_epoch_end(self,
  training_step_outputs) 函数；
• ● validation_step 后面都紧跟着其相应的
  validation_step_end(self，batch_parts)和validation_epoch_end(self,
  training_step_outputs) 函数；
• ● test_step 后面都紧跟着其相应的 test_step_end(self，batch_parts)和
  test_epoch_end(self, training_step_outputs) 函数；

每个主要函数后面跟一个step_end和epoch_end
以训练函数为例：

def training_step(self, batch, batch_idx):
    x, y = batch
    y_hat = self.model(x)
    loss = F.cross_entropy(y_hat, y)
    pred = ...
    return {'loss': loss, 'pred': pred}

def training_step_end(self, batch_parts):
    '''
    当gpus=0 or 1时，这里的batch_parts即为traing_step的返回值（已验证）
    当gpus>1时，这里的batch_parts为list，list中每个为training_step返回值，list[i]为i号gpu的返回值（这里未验证）
    '''
    gpu_0_prediction = batch_parts[0]['pred']
    gpu_1_prediction = batch_parts[1]['pred']

    # do something with both outputs
    return (batch_parts[0]['loss'] + batch_parts[1]['loss']) / 2

def training_epoch_end(self, training_step_outputs):
    '''
    当gpu=0 or 1时，training_step_outputs为list，长度为steps的数量（不包括validation的步数，当你训练时，你会发现返回list<训练时的steps数，这是因为训练时显示的steps数据还包括了validation的，若将limit_val_batches=0.，即关闭validation，则显示的steps会与training_step_outputs的长度相同）。list中的每个值为字典类型，字典中会存有`training_step_end()`返回的键值，键名为`training_step()`函数返回的变量名，另外还有该值是在哪台设备上(哪张GPU上)，例如{device='cuda:0'}
    '''
    for out in training_step_outputs:
       # do something with preds

代码分析：
training_step(self, batch, batch_idx)，类似于pytorch中的forward函数，解决的是预测与损失计算，并返回预测与损失结果。

training_step_end(self, batch_parts)，方法是用于对每个训练批次结束后进行一些操作的生命周期钩子。它允许在单个训练批次完成后进行额外的处理，例如收集和记录训练统计信息、生成日志、或执行其他自定义操作。

training_epoch_end(self, training_step_outputs)，每个训练轮次（epoch）结束时进行一些操作。它允许您在一个轮次结束后对训练过程进行一些全局性的处理，例如收集并汇总各个批次的信息，生成日志，计算额外的指标，或进行其他定制的操作。

import pytorch_lightning as pl
class MyLightningModule(pl.LightningModule):
    def __init__(self):
        super().__init__()

    def training_epoch_end(self, training_step_outputs):
        # 在轮次结束时，收集和汇总信息
        loss_values = [output['loss'] for output in training_step_outputs]
        avg_loss = sum(loss_values) / len(loss_values)

        # 记录日志
        self.log('train_loss', avg_loss)

        # 可以执行其他自定义操作，如更新学习率或保存模型检查点
        # self.lr_scheduler.step()

        # 可以返回一些信息，这些信息将被记录到 TensorBoard 中
        return {'epoch_loss': avg_loss}

1、Train

训练的主要方法是重写def training_setp(self, batch, batch_idx)函数，并返回要反向传播的loss
batch代表dataloader加载的数据，batch_idx对应每个batch的索引

def training_setp(self, batch, batch_idx):
    image, label = batch
    pred = self.forward(iamge)
    loss = ...
    # 一定要返回loss
    return loss

2、Validation(验证频率)

在pytorch中，需要自己完成训练，验证，测试的代码编写，但在pytorch Lightning框架中，验证可以通过一句话配置check_val_every_n_epoch，默认1轮验证一次，也可以通过下面的内容设置。
trainer = Trainer(check_val_every_n_epoch=5)
当一个epoch过大时，希望在epoch训练一部分就进行验证的情况，该框架也给出了相应的解决方法val_check_interval

# 单个epoch进行 50% 调用校验函数一次，传入float
trainer = Trainer(val_check_interval=0.5)
#也可以训练完多少个batch后调用一次校验函数，是传入int型
trainer = Trainer(val_check_interval=100) # 每训练100个batch校验一次

3、test

训练过程中不调用test，只会执行train和val两部分函数，因此如果需要保存一些信息，就要放到validation中

# 获取恢复了权重和超参数等的模型
model = MODEL.load_from_checkpoint(checkpoint_path='my_model_path/heiheihei.ckpt')
# 修改测试时需要的参数，例如预测的步数等
model.pred_step = 1000
# 定义trainer, 其中limit_test_batches表示取测试集中的0.05的数据来做测试
trainer = pl.Trainer(gpus=1, precision=16, limit_test_batches=0.05)
# 测试，自动调用test_step(), 其中dm为数据集，放在下面讲
trainer.test(model=dck, datamodule=dm)

参考文献：
https://blog.csdn.net/u013250861/article/details/123591130
GPT3.5相关内容