跟踪训练
指标跟踪
模型训练时,通常需要查看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)