PyTorch中神经网络、损失函数与优化器的关联机制解惑

Great question—this is one of the key "magic" parts of PyTorch (and other autograd frameworks) that trips up a lot of folks at first. Let me break down how all these pieces connect, step by step:

The Hidden Link: The Computation Graph #

When you run outputs = net(inputs), PyTorch doesn’t just compute the output tensor—it builds a dynamic computation graph behind the scenes. This graph tracks every operation that transforms the input into the output, and crucially, it keeps references to all the network parameters (weights, biases) involved in those operations. Every parameter in your nn.Module has a requires_grad=True flag by default, which tells PyTorch to track gradients for that parameter.

Loss Ties Into the Graph #

When you calculate loss = criterion(outputs, labels), the loss tensor isn’t just a scalar value—it’s the final node in that computation graph. Even though you’re only passing outputs and labels to the loss function, the loss tensor retains a direct connection to the entire chain of operations (and thus the network parameters) that produced the outputs.

loss.backward() Fills in the Gradients #

This is the critical step that bridges the gap between loss and parameters. Calling backward() tells PyTorch to traverse the computation graph in reverse, starting from the loss tensor. For every parameter in the graph, it calculates how much that parameter contributes to the loss (the gradient) and stores that value in the parameter’s .grad attribute.

After running loss.backward(), every weight and bias in your net will have a .grad value that represents how changing that parameter would affect the loss.

Optimizer Uses Stored Gradients to Update Parameters #

When you initialized the optimizer with net.parameters(), you gave it direct references to all those parameter tensors. The optimizer doesn’t need to "communicate" with the loss function directly—by the time you call optimizer.step(), each parameter already has its gradient stored in .grad.

The optimizer simply reads those gradients and applies its update rule (like SGD, Adam, etc.) to adjust the parameter values in a way that reduces the loss.

Concrete Example to Illustrate #

Let’s use a tiny network to make this tangible:

import torch
import torch.nn as nn
import torch.optim as optim

# Tiny network: single linear layer (y = wx + b)
net = nn.Linear(1, 1)
criterion = nn.MSELoss()
optimizer = optim.SGD(net.parameters(), lr=0.01)

# Dummy training data
inputs = torch.tensor([[1.0]])
labels = torch.tensor([[2.0]])

# Forward pass: build computation graph
outputs = net(inputs)
loss = criterion(outputs, labels)

# Check gradients before backward pass (they're None)
print("Before backward: w.grad =", net.weight.grad)

# Trigger backward pass to compute gradients
loss.backward()
# Now gradients are stored in the parameters
print("After backward: w.grad =", net.weight.grad)
print("After backward: b.grad =", net.bias.grad)

# Optimizer updates parameters using stored gradients
optimizer.step()
# Now w and b are adjusted to reduce the loss

Recap #

To sum it up:

• The computation graph creates the hidden connection between network parameters, outputs, and loss.
• loss.backward() populates the gradient values for each parameter.
• The optimizer uses these precomputed gradients to update the network weights.

No direct communication between the loss function and optimizer is needed—PyTorch’s autograd system handles all the wiring behind the scenes.

内容的提问来源于stack exchange，提问作者momo