Training - PyTorch Lightning 分布式训练的 global_step 参数 (accumulate_grad_batches)

SpikeKing2024-04-18 15:45

欢迎关注我的CSDN：https://spike.blog.csdn.net/

本文地址：https://blog.csdn.net/caroline_wendy/article/details/137640653

在 PyTorch Lightning 中，pl.Trainer 的 accumulate_grad_batches 参数允许在执行反向传播和优化器步骤之前，累积多个批次的梯度。这样，可以增加有效的批次大小，而不会增加内存开销。例如，如果设置 accumulate_grad_batches=8，则会在执行优化器的 .step() 方法之前，累积 8 个批次的梯度。

accumulate_grad_batches 与 global_step 的关系：

global_step 会在每次调用优化器的 .step() 方法后递增。
使用梯度累积，global_step 增长小于批次(batch) 的数量
多个批次贡献到 1 个 global_step 的更新中。

例如，如果 accumulate_grad_batches=8，那么每 8 个批次，只会增加 1 次 global_step，如果多卡，则 global_step 表示单卡的次数。日志，如下：

bash 复制代码

[INFO] [CL] global_step: 0, iter_step: 8
[INFO] [CL] global_step: 1, iter_step: 16

其中 pl.Trainer 的源码：

bash 复制代码

    trainer = pl.Trainer(
        accelerator="gpu",
        # ...
        accumulate_grad_batches=args.accumulate_grad,
        strategy=strategy,  # 多机多卡配置
        num_nodes=args.num_nodes,  # 节点数
        devices=1,  # 每个节点 GPU 卡数
    )

输出日志：

bash 复制代码

log = {'epoch': self.trainer.current_epoch, 'step': self.trainer.global_step}
wandb.log(log)