使用PyTorch Lightning与W&B时训练未完成静默崩溃问题求助

问题详情 #

• 使用PyTorch Lightning结合wandb开展训练实验，预设训练轮次为1000轮，但训练执行到第999轮时，wandb提示进程已完成，实验未跑完全部轮次，且wandb界面显示实验状态为crashed
• 原有代码执行流程：
  • 通过wandb_logger.watch()监控模型参数
  • 完成数据准备后，启动pl.Trainer开始训练
  • 训练结束后调用wandb_logger.experiment.finish()
  • 需按顺序执行多个实验，每个实验末尾都会调用finish()

解决方法 #

在调用finish()之前，添加代码手动标记实验状态为成功：

wandb_logger.experiment.finalize('success')
wandb_logger.experiment.finish()

添加后实验可正常跑完1000轮，wandb界面不再标记为crashed。

内容的提问来源于stack exchange，提问作者James Arten