理论
功能:将梯度清零。
由于 PyTorch 不会自动清零梯度,所以在每一次更新前会进行此操作。
案例
# coding: utf-8
import torch
import torch.optim as optim
# ----------------------------------- zero_grad
w1 = torch.randn(2, 2)
w1.requires_grad = True
w2 = torch.randn(2, 2)
w2.requires_grad = True
optimizer = optim.SGD([w1, w2], lr=0.001, momentum=0.9)
# optimizer.param_groups:{'params': [tensor([[0.3673, 0.6950],[0.4002, 0.1766]],requires_grad=True)
# 参数组:optimizer.param_groups 参数组第一个参数:optimizer.param_groups[0]
# 参数组第一个参数中params键值:optimizer.param_groups[0]['params']
# 参数组第一个参数中params键值对应的第一个元素:optimizer.param_groups[0]['params'][0]
optimizer.param_groups[0]['params'][0].grad = torch.randn(2, 2)
print('参数w1的梯度:')
print(optimizer.param_groups[0]['params'][0].grad, '\n') # 参数组,第一个参数(w1)的梯度
optimizer.zero_grad()
print('执行zero_grad()之后,参数w1的梯度:')
print(optimizer.param_groups[0]['params'][0].grad) # 参数组,第一个参数(w1)的梯度
结果
参数w1的梯度:
tensor([[ 0.5045, 0.7872],
[ 0.4429, -0.2255]])
执行zero_grad()之后,参数w1的梯度:
tensor([[0., 0.],
[0., 0.]])