在反向传播计算完所有参数的梯度后,还需要使用优化方法更新网络的权重和参数。例如,随机梯度下降法(SGD)的更新策略如下:
weight = weight - learning_rate * gradient
手动实现如下:
learning_rate = 0.01
for f in net.parameters():
f.data.sub_(f.grad.data * learning_rate) # inplace减法
在torch.optim
中实现了深度学习中绝大多数的优化方法,如RMSProp、Adam、SGD等。
import torch.optim as optim
# 新建一个优化器,指定要调整的参数和学习率
optimizer = optim.SGD(net.parameters(), lr = 0.01)
# 在训练过程中先将梯度清零(与net.zero_grad()效果一样)
optimizer.zero_grad()
# 计算损失
output = net(input)
loss = criterion(output, target)
# 反向传播
loss.backward()
# 更新参数
optimizer.step()