举个例子说明optimizer.zero_grad()与optimizer.step()的用法
时间: 2023-08-07 19:11:07 浏览: 42
假设我们有一个神经网络模型和一个优化器对象,可以这样使用optimizer.zero_grad()和optimizer.step():
```python
import torch
import torch.nn as nn
import torch.optim as optim
# 创建一个简单的神经网络模型
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(10, 5)
self.fc2 = nn.Linear(5, 2)
def forward(self, x):
x = self.fc1(x)
x = self.fc2(x)
return x
# 实例化模型和优化器对象
net = Net()
optimizer = optim.SGD(net.parameters(), lr=0.01)
# 定义一个输入和对应的标签
input_data = torch.randn(1, 10)
target = torch.Tensor([1, 0])
# 计算模型输出和损失值
output = net(input_data)
criterion = nn.MSELoss()
loss = criterion(output, target)
# 梯度清零,反向传播,更新参数
optimizer.zero_grad() # 将所有参数的梯度置零
loss.backward() # 反向传播,计算参数梯度
optimizer.step() # 更新参数
```
在上述代码中,我们首先定义了一个简单的神经网络模型和一个优化器对象。接着,我们定义了一个输入和对应的标签,并计算了模型的输出和损失值。然后,我们使用optimizer.zero_grad()将所有参数的梯度清零,使用loss.backward()进行反向传播计算参数梯度,最后使用optimizer.step()更新参数。这些步骤通常在训练神经网络时会反复执行多次,以不断优化模型的性能。