LightningModule 类的self.log方法介绍

1️⃣ self.log 方法简介

self.log 是 PyTorch Lightning 提供的方法，主要用于 记录训练、验证、测试和推理过程中的指标（如 loss、accuracy 等）。
它支持 TensorBoard、WandB、CSV Logger 等自动日志记录方式，无需手动管理 writer.add_scalar()。

2️⃣ self.log 方法的参数

self.log(
    name: str,  # 记录的名称 (e.g., "train/loss")
    value: Union[torch.Tensor, float, int],  # 记录的值 (loss, accuracy 等)
    prog_bar: bool = False,  # 是否在进度条中显示
    on_step: bool = False,  # 是否在 step 级别记录
    on_epoch: bool = False,  # 是否在 epoch 级别记录
    logger: bool = True,  # 是否记录到 logger (如 TensorBoard)
    sync_dist: bool = False,  # 在多 GPU 训练时是否同步
    reduce_fx: Callable = torch.mean,  # 多 GPU 情况下的归约函数 (默认取平均)
)

📌 关键参数解析

参数	作用
name	记录指标的名称 (e.g., "train/loss"、"val/accuracy")
value	需要记录的值，通常是 torch.Tensor 或 float
prog_bar	是否在 进度条 显示 (默认 False)
on_step	是否在 每个 step 记录 (适用于训练)
on_epoch	是否在 每个 epoch 结束后 记录 (适用于验证)
logger	是否写入到日志 (默认 True)
sync_dist	多 GPU 训练时，是否同步日志数据
reduce_fx	多 GPU 训练时的归约方式 (默认 torch.mean)

---

3️⃣ self.log 的使用示例

🌟 示例 1：记录训练损失

import torch
import pytorch_lightning as pl

class MyModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = torch.nn.Linear(10, 1)
        self.loss_fn = torch.nn.MSELoss()

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.model(x)
        loss = self.loss_fn(y_hat, y)

        # 记录 loss
        self.log("train/loss", loss, prog_bar=True, on_step=True, on_epoch=True)
        return loss

📌 解释：

• train/loss 记录 损失值。
• prog_bar=True：在 进度条 中显示 loss。
• on_step=True：在 每个 batch 训练 step 记录 loss。
• on_epoch=True：在 每个 epoch 结束后 计算 loss 均值。

---

🌟 示例 2：记录验证准确率

def validation_step(self, batch, batch_idx):
    x, y = batch
    y_hat = self.model(x)
    acc = (y_hat.argmax(dim=1) == y).float().mean()

    # 记录准确率
    self.log("val/accuracy", acc, prog_bar=True, on_step=False, on_epoch=True)

📌 解释：

• 只在 epoch 结束时 记录准确率，不在 step 级别记录 (on_step=False, on_epoch=True)。

---

🌟 示例 3：记录多个指标

def training_step(self, batch, batch_idx):
    x, y = batch
    y_hat = self.model(x)
    loss = self.loss_fn(y_hat, y)
    acc = (y_hat.argmax(dim=1) == y).float().mean()

    self.log("train/loss", loss, prog_bar=True, on_step=True, on_epoch=True)
    self.log("train/accuracy", acc, prog_bar=True, on_step=False, on_epoch=True)

📌 解释：

• train/loss 每个 step 和 epoch 记录，并显示在进度条 (prog_bar=True)。
• train/accuracy 仅在 epoch 结束后记录 (on_step=False, on_epoch=True)。

---

🌟 示例 4：多 GPU 训练时日志同步

def training_step(self, batch, batch_idx):
    x, y = batch
    y_hat = self.model(x)
    loss = self.loss_fn(y_hat, y)

    # 在多 GPU 模式下，同步 loss
    self.log("train/loss", loss, on_step=True, on_epoch=True, sync_dist=True)
    return loss

📌 解释：

• sync_dist=True：在 分布式训练 (DDP) 时，同步 loss 计算结果。

---

4️⃣ self.log vs self.log_dict

• self.log：记录单个指标
• self.log_dict：同时记录多个指标

def training_step(self, batch, batch_idx):
    x, y = batch
    y_hat = self.model(x)
    loss = self.loss_fn(y_hat, y)
    acc = (y_hat.argmax(dim=1) == y).float().mean()

    # 一次性记录多个指标
    self.log_dict({"train/loss": loss, "train/accuracy": acc}, on_step=True, on_epoch=True)

🔹 总结

✅ self.log 是 PyTorch Lightning 提供的 自动日志记录方法，无需手动管理 TensorBoard。
✅ 支持 step 和 epoch 级别的日志记录。
✅ 适用于 训练 (train)、验证 (val)、测试 (test) 和推理 (predict) 阶段的指标记录。
✅ 在 多 GPU 训练 时，使用 sync_dist=True 确保日志同步。

🚀 self.log 让日志记录更简洁高效，适用于 AlphaFold 训练中的 loss、LDDT、GDT-TS 等指标记录！