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

相关推荐
codeGoogle8 分钟前
大厂研发之谜:千亿投入砸出利润大缩水
前端·人工智能·后端
PP东17 分钟前
Pyhton基础之多继承、多态
开发语言·python
豆浩宇25 分钟前
Conda环境隔离和PyCharm配置,完美同时运行PaddlePaddle和PyTorch
人工智能·pytorch·算法·计算机视觉·pycharm·conda·paddlepaddle
京东云开发者34 分钟前
DeepSeek冲击(含本地化部署实践)
人工智能
@国境以南,太阳以西40 分钟前
基于Grad-CAM(Gradient-weighted Class Activation Mapping)的可解释性分析
人工智能·深度学习
菜鸟的日志1 小时前
【音频字幕】构建一个离线视频字幕生成系统:使用 WhisperX 和 Faster-Whisper 的 Python 实现
python·whisper·音视频
AI人工智能+1 小时前
表格识别技术:通过计算机视觉和OCR,实现非结构化表格向结构化数据的转换,推动数字化转型。
人工智能·计算机视觉·ocr
小宁爱Python1 小时前
基于 Django+Vue3 的 AI 海报生成平台开发(海报模块专项)
人工智能·python·django
破烂儿1 小时前
基于机器学习的缓存准入策略研究
人工智能·机器学习·缓存
红豆怪怪2 小时前
[LeetCode 热题 100] 32. 最长有效括号
数据结构·python·算法·leetcode·动态规划·代理模式