pytorch-lightning的trainer的optimaze_step的详细分析

wzy-666

已于 2024-01-14 11:02:25 修改

阅读量522

点赞数 1

文章标签： pytorch 人工智能 python

于 2023-09-15 10:56:51 首次发布

本文链接：https://blog.youkuaiyun.com/C_C666/article/details/132897036

版权

本文概述了在使用PyTorchLightning时，如何在训练过程中管理优化器步骤，包括梯度缩放（如scaler.scale()和scaler.step()），梯度裁剪，以及处理NaN梯度的策略。重点介绍了`on_train_batch_start`到`on_train_batch_end`的钩子函数在优化过程中的作用。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

on_train_batch_start
optimaze_step
- training_step
- on_before_zero_grad
- scaler.scale(loss) 将loss放大
- on_before_backward
- model.backward 梯度反传
- on_after_backward
- scaler.unscale_(optimizer) 将grad还原
- on_before_optimizer_step
- _clip_gradients梯度裁剪
- scaler.step(optimizer) 如果发现nan梯度，则optimizer跳过
- scaler.update() 更新缩放器
on_train_batch_end