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() 来根据梯度更新网络参数的值。

相关推荐
ai_xiaogui2 分钟前
一键部署AI工具!用AIStarter快速安装ComfyUI与Stable Diffusion
人工智能·stable diffusion·部署ai工具·ai应用市场教程·sd快速部署·comfyui一键安装
Tipriest_8 分钟前
Python关键字梳理
python·关键字·keyword
聚客AI1 小时前
Embedding进化论:从Word2Vec到OpenAI三代模型技术跃迁
人工智能·llm·掘金·日新计划
weixin_387545641 小时前
深入解析 AI Gateway:新一代智能流量控制中枢
人工智能·gateway
im_AMBER1 小时前
学习日志05 python
python·学习
大虫小呓2 小时前
Python 处理 Excel 数据 pandas 和 openpyxl 哪家强?
python·pandas
聽雨2372 小时前
03每日简报20250705
人工智能·社交电子·娱乐·传媒·媒体
哪 吒2 小时前
2025B卷 - 华为OD机试七日集训第5期 - 按算法分类,由易到难,循序渐进,玩转OD(Python/JS/C/C++)
python·算法·华为od·华为od机试·2025b卷
二川bro2 小时前
飞算智造JavaAI:智能编程革命——AI重构Java开发新范式
java·人工智能·重构
acstdm2 小时前
DAY 48 CBAM注意力
人工智能·深度学习·机器学习