PyTorch中部分ADAM配置训练单步耗时递增，SGD无此问题求助

针对你遇到的PyTorch Lightning 1.8.4中ADAM优化器单步训练耗时随时间递增（特定种子下）、SGD无此现象的问题，给出以下排查方向和解决建议：

1. 检查数值稳定性与优化器状态异常 #

ADAM依赖一阶动量（exp_avg）和二阶动量（exp_avg_sq）的累积，某些随机种子可能导致训练过程中这些状态张量出现NaN/Inf或者异常大的数值，进而增加后续step的计算开销。

• 排查方法：在opt.step()后，打印优化器的状态字典（opt.state_dict()），检查exp_avg、exp_avg_sq是否有异常值；同时监控模型参数、梯度的数值范围。
• 解决方法：
  • 添加梯度裁剪：在manual_backward后、opt.step()前执行梯度裁剪，限制梯度大小避免数值爆炸：

    self.clip_gradients(opt, gradient_clip_val=1.0, gradient_clip_algorithm='norm')
    

  • 启用PyTorch的自动数值检查：在训练前添加torch.autograd.set_detect_anomaly(True)，捕获异常的梯度计算节点。

2. 修正手动优化模式下的细节问题 #

禁用自动优化后，原生PyTorch优化器的使用可能忽略了PyTorch Lightning的内部状态管理，导致潜在的计算图残留或内存累积：

• 排查方法：对比使用use_pl_optimizer=True的情况，看耗时是否依然递增；监控GPU/CPU内存占用随训练步数的变化。
• 解决方法：
  • 将opt = self.optimizers(use_pl_optimizer=False)改为use_pl_optimizer=True，使用PL封装的优化器，确保状态被正确跟踪。
  • 确保log_dict中的张量脱离计算图：当前log_dict['time_step']依赖loss的计算图，会被PL保留导致内存累积，修改为：

    log_dict['time_step'] = torch.tensor(time.time() - tic).detach()
    

3. 替换ADAM实现或调整优化器参数 #

原生ADAM的权重衰减实现可能在某些场景下引发额外计算开销，尝试使用AdamW优化器替代，或调整参数：

• 解决方法：
  • 将优化器替换为AdamW（PyTorch 1.13.1支持），修改优化器初始化代码：

    from torch.optim import AdamW
    def configure_optimizers(self):
        return AdamW(self.parameters(), lr=1e-3)
    

  • 暂时将weight_decay设为0，测试耗时是否依然递增，排查是否是权重衰减导致的计算开销累积。

4. 排查内存泄漏 #

手动优化模式下，未正确释放的张量或计算图会导致内存占用逐步上升，间接增加单步训练耗时：

• 解决方法：
  • 在training_step末尾添加内存清理：如果使用GPU，执行torch.cuda.empty_cache()；如果是CPU，执行del loss后调用gc.collect()。
  • 使用PyTorch的内存分析工具（如torch.utils.tensorboard.summary.record_gradients）或第三方工具（如pytorch-memlab）定位内存泄漏点。

内容的提问来源于stack exchange，提问作者Pablo Sanchez