PyTorch中的model.zero_grad() vs op

文章作者：Tyan
博客：noahsnail.com  |  CSDN  |  简书

1. 引言

在PyTorch中，对模型参数的梯度置0时通常使用两种方式：model.zero_grad()和optimizer.zero_grad()。二者在训练代码都很常见，那么二者的区别在哪里呢？

2. model.zero_grad()

model.zero_grad()的作用是将所有模型参数的梯度置为0。其源码如下：

for p in self.parameters():
    if p.grad is not None:
        p.grad.detach_()
        p.grad.zero_()

3. optimizer.zero_grad()

optimizer.zero_grad()的作用是清除所有优化的torch.Tensor的梯度。其源码如下：

for group in self.param_groups:
    for p in group['params']:
        if p.grad is not None:
            p.grad.detach_()
            p.grad.zero_()

4. 总结

• 当使用optimizer = optim.Optimizer(net.parameters())设置优化器时，此时优化器中的param_groups等于模型中的parameters()，此时，二者是等效的，从二者的源码中也可以看出来。

• 当多个模型使用同一个优化器时，二者是不同的，此时需要根据实际情况选择梯度的清除方式。

• 当一个模型使用多个优化器时，二者是不同的，此时需要根据实际情况选择梯度的清除方式。

References

1. https://discuss.pytorch.org/t/model-zero-grad-or-optimizer-zero-grad/28426

2. https://discuss.pytorch.org/t/whats-the-difference-between-optimizer-zero-grad-vs-nn-module-zero-grad/59233

3. https://discuss.pytorch.org/t/zero-grad-optimizer-or-net/1887

4. https://pytorch.org/docs/stable/optim.html

5. https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module