Pytorch:optimizer.zero_grad(), loss.backward(), optimizer.step()

在训练过程中常看到如下代码:

python 复制代码
optimizer.zero_grad(set_to_none=True)
grad_scaler.scale(loss).backward()
grad_scaler.step(optimizer)
grad_scaler.update()

这三个函数的作用是:

  • 在训练过程中先调用 optimizer.zero_grad() 清空梯度
  • 再调用 loss.backward() 反向传播
  • 最后调用optimizer.step()更新模型参数:

optimizer.zero_grad():

Optimizer类在实例化时会在构造函数中创建一个param_groups列表

param_group['params']:由传入的模型参数组成的列表,即实例化Optimizer类时传入该group的参数,如果参数没有分组,则为整个模型的参数model.parameters(),每个参数是一个torch.nn.parameter.Parameter对象。

optimizer.zero_grad()函数会遍历模型的所有参数,通过:

  • p.grad.detach_()方法截断反向传播的梯度流;
  • 再通过p.grad.zero_()函数将每个参数的梯度值设为0,即上一次的梯度记录被清空。

具体来说,optimizer.zero_grad() 会将优化器中所有可学习参数的梯度设为 0。这样,在下一次前向传递计算和反向传播计算时,之前的梯度就不会对当前的梯度产生影响。

在反向传播计算梯度之前对上一次迭代时记录的梯度清零,参数 set_to_none 设置为 True 时会直接将参数梯度设置为 None,从而减小内存使用,但通常情况下不建议设置这个参数,因为梯度设置为 None 和 0 在 PyTorch 中处理逻辑会不一样。

python 复制代码
def zero_grad(self, set_to_none: bool = False):
        r"""Sets the gradients of all optimized :class:`torch.Tensor` s to zero.

        Arguments:
            set_to_none (bool): instead of setting to zero, set the grads to None.
                This is will in general have lower memory footprint, and can modestly improve performance.
                However, it changes certain behaviors. For example:
                1. When the user tries to access a gradient and perform manual ops on it,
                a None attribute or a Tensor full of 0s will behave differently.
                2. If the user requests ``zero_grad(set_to_none=True)`` followed by a backward pass, ``.grad``s
                are guaranteed to be None for params that did not receive a gradient.
                3. ``torch.optim`` optimizers have a different behavior if the gradient is 0 or None
                (in one case it does the step with a gradient of 0 and in the other it skips
                the step altogether).
        """
        for group in self.param_groups:
            for p in group['params']:
                if p.grad is not None:
                    if set_to_none:
                        p.grad = None
                    else:
                        if p.grad.grad_fn is not None:
                            p.grad.detach_()
                        else:
                            p.grad.requires_grad_(False)
                        p.grad.zero_()

因为训练的过程通常使用mini-batch方法,所以如果不将梯度清零的话,梯度会与上一个batch的数据相关,因此该函数要写在反向传播和梯度下降之前。

更多详细内容,参考:Pytorch:torch.optim模块

loss.backward():

PyTorch的反向传播(即tensor.backward())是通过autograd包来实现的,autograd包会根据tensor进行过的数学运算来自动计算其对应的梯度。

具体来说,torch.tensor是autograd包的基础类,如果你设置tensor的requires_grads为True,就会开始跟踪这个tensor上面的所有运算,如果你做完运算后使用tensor.backward(),所有的梯度就会自动运算,tensor的梯度将会累加到它的.grad属性里面去。

具体参考:Pytorch:backward()函数详解

optimizer.step()

optimizer.step():此方法主要完成一次模型参数的更新。

当使用 backward() 计算网络参数的梯度后,需要使用 optimizer.step() 来根据梯度更新网络参数的值。

相关推荐
芯盾时代18 分钟前
安全大模型智驱网络和数据安全效能跃迁
网络·人工智能·安全·网络安全
彩讯股份3006341 小时前
打造多模态交互新范式|彩讯股份中标2025年中国移动和留言平台AI智能体研发项目
人工智能
思通数科大数据舆情1 小时前
工业安全零事故的智能守护者:一体化AI智能安防平台
人工智能·安全·目标检测·计算机视觉·目标跟踪·数据挖掘·知识图谱
AI360labs_atyun2 小时前
2025 高考:AI 都在哪些地方发挥了作用
人工智能·科技·ai·高考
Yxh181377845543 小时前
短视频矩阵系统技术saas源头6年开发构架
人工智能·矩阵
m0_634448893 小时前
图上合成:用于大型语言模型持续预训练的知识合成数据生成
人工智能·语言模型·自然语言处理
亚林瓜子3 小时前
AWS中国云的定时任务(AWS EventBridge+AWS Lambda)
python·云计算·aws·lambda·定时任务·event·cron
摆渡搜不到你4 小时前
某腾X视频下载器2.1
python·音视频
科雷软件测试4 小时前
Python格式化:让数据输出更优雅
python
非小号4 小时前
多模态分类案例实现
python·分类·数据挖掘