仅计算DNN前两层参数时反向传播为何更高效？

核心原因是PyTorch的自动微分引擎会跳过不必要的计算步骤与内存存储，具体可以从以下几点拆解：

1. 跳过非目标参数的梯度计算 #

反向传播的本质是通过链式法则计算损失对每个参数的偏导数。当你只需要前两层参数的梯度时，后续所有层的参数偏导计算会被直接跳过：

• 比如网络结构是输入→层1→层2→层3→层4→损失，原本需要计算损失对层4、层3、层2、层1所有参数的偏导；
• 现在仅需层2、层1的参数梯度，那么层4、层3的参数偏导计算步骤会被完全省略，这部分计算量的节省在深层网络中会非常显著。

2. 减少中间张量的存储与加载开销 #

前向传播时，PyTorch会存储中间激活值（用于反向传播的链式法则计算）。如果后续层的参数被设置为requires_grad=False，或者通过torch.autograd.grad指定仅计算前两层参数的梯度，PyTorch会自动跳过后续层中间激活的梯度信息存储：

• 这不仅节省了内存占用，还减少了反向传播时加载这些张量的额外开销，进一步提升了计算效率。

3. 精准追踪必要的计算路径 #

当使用torch.autograd.grad直接指定目标参数时，自动微分引擎会只追踪从损失到这些参数的计算路径，对路径外的节点（比如后续层的参数和无关中间激活）不进行任何反向传播处理，相当于只完成了反向传播中“必要的部分”，而非全链路计算。

举个PyTorch实操例子 #

假设我们有一个4层全连接网络，仅需计算前两层的梯度：

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 20)
        self.fc2 = nn.Linear(20, 30)
        self.fc3 = nn.Linear(30, 10)
        self.fc4 = nn.Linear(10, 1)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = torch.relu(self.fc3(x))
        return self.fc4(x)

# 方式1：关闭后续层的梯度追踪
net = Net()
for param in list(net.fc3.parameters()) + list(net.fc4.parameters()):
    param.requires_grad = False

x = torch.randn(32, 10)
loss = net(x).sum()
loss.backward()
# 仅fc1、fc2的参数有grad值，fc3、fc4的grad为None

# 方式2：直接指定计算目标参数的梯度
net = Net()
x = torch.randn(32, 10)
loss = net(x).sum()
grads = torch.autograd.grad(loss, list(net.fc1.parameters()) + list(net.fc2.parameters()))

这两种方式都能避免后续层的冗余计算，相比全网络梯度计算，在参数规模大、层数多的场景下，效率提升会非常明显。

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