如何使用PyTorch的zero_grad()、backward()和step()

Question

我有一个由 nn.module 创建的基本线性回归类，这是该类：

class LinearRegressionModel2(nn.Module):
  def __init__(self):
    super().__init__()
    # Use nn.Linear() for creating the model parameters (also called linear transform, probing layer, fully connected layer, dense layer)
    self.linear_layer = nn.Linear(in_features = 1,
                                  out_features = 1)

  def forward(self, x: torch.Tensor) -> torch.Tensor:
    return self.linear_layer(x)

我尝试在循环步骤之前使用测试和训练循环进行基本预测，我创建了损失函数和优化器，以下是相关代码：

torch.manual_seed(42)
model_1 = LinearRegressionModel2()

# Setup Loss Function
loss_fn = nn.L1Loss() # Same ass MAE

# Setup our optimizer
optimizer = torch.optim.SGD(params = model_1.parameters(),
                            lr = 0.01, )

epochs = 200

for epoch in range(epochs):
  model_1.train()

  # 1. Forward pass
  y_pred = model_1(X_train)

  # 2. Calculate the loss
  train_loss = loss_fn(y_pred, y_train)
  
  # 3. Optimizer zero grad
  optimizer.zero_grad()

  # 4. Perform backpropagation
  train_loss.backward()

  # 5. Optimizer step
  optimizer.step()

  ### Testing
  model_1.eval()
  with torch.inference_mode():
    test_pred = model_1(X_test)
    test_loss = loss_fn(test_pred, y_test)

  # Print out whats happening

  if epoch % 10 == 0:
    print(f"Epoch: {epoch} | Train Loss: {train_loss} | Test Loss: {test_loss}")

但是我无法理解4.和5.步骤，当我在网上搜索时，我发现

optimizer.zero_grad()

用于重置每批的梯度步骤。 3. 步骤还可以，但是在第 4. 步骤中如何使用 back() 仅使用数字，以及在第 4. 步骤之后优化器如何知道损失 train_loss.backward() 以及这两个步骤如何协同工作，因为有代码中没有任何联系。总之，第 3. 4. 和 5. 步骤如何协同工作？

我在谷歌中搜索并询问了chatgpt，但我不明白其中的关系以及它们是如何工作的/

Answer 1

4. 步骤如何使用数字来向后（）工作

您的损失函数是一个凸函数，在某种程度上测量真实值和预测之间的误差。给定该值、输入以及导致该值的所有计算，您可以计算每个操作的导数，例如网络的密集层。

这两个步骤如何协同工作，因为没有任何 代码中的连接。

初始化优化器时，您必须提供所有需要优化的参数：

optimizer = torch.optim.SGD(params = **model_1.parameters()**,
                            lr = 0.01, )

在内部，pytorch 对所有发生的计算执行所有必需的魔法和簿记。

所以，总的来说：

1. 给定一些输入，当前的参数集，通过前向传播计算预测
2. 使用损失函数计算预测值与真实值之间的差异
3. 使用您在步骤 2 中测量的误差（向后传递）计算模型每个参数的导数。这一步是pytorch内部执行的，你不应该自己做任何事情
4. 使用一些规则更新模型的参数 - optimizationr.step()
5. 在下一步 - Zero_grad() 之前删除所有结果，以便您可以重复此过程