6.pytorch lightning 跟踪训练

跟踪训练

指标跟踪

模型训练时，通常需要查看trainig_loss确认训练过程是否正常，以及根据验证集评估的精度保存模型，则需要将这些信息记录，用于logger自动绘制学习曲线等。另外，你可以需要查看模型推理结果是否正常，则可以记录模型预测的图片、文本等结果。训练管理器Trainer会根据相应的设置（比如ModelCheckpoint 保存模型的callback）执行处理流程。

为了跟踪训练，在 LightningModule的方法中调用self.log /self.log_dict。

class LitModel(pl.LightningModule):
    def training_step(self, batch, batch_idx):
        # 记录单个标题
        value = ...
        self.log("some_value", value)

        # 记录字典项
        values = {"loss": loss, "acc": acc, "metric_n": metric_n}  # add more items if needed
self.log_dict(values)

可视化

使用logger 工具记录图片，上传绘图等，支持的工具有tensorboard, wandb, comet, mlflow, Neptune。以本地的tensorboard为例，创建Trainer时指定logger参数：

from lightning.pytorch.loggers import TensorBoardLogger

logger = TensorBoardLogger()
trainer = Trainer(logger=logger)

然后在LightningModule的模型中调用tensorboard工具内的方法

class LitModel(LightningModule):
    def any_lightning_module_function_or_hook(self):
        tensorboard_logger = self.logger.experiment# 获取tensorboard实例
        fake_images = torch.Tensor(32, 3, 28, 28)
        tensorboard_logger.add_image("generated_images", fake_images, 0)