前向传播使用张量索引取值时PyTorch如何执行反向微分

部分代码改编自《Deep learning with Pytorch》一书

线性回归脚本说明 #

目标是根据t_u预测t_c，代码如下：

import torch

t_c = torch.tensor([0.5, 14.0, 15.0, 28.0, 11.0, 8.0,
                    3.0, -4.0, 6.0, 13.0, 21.0])
t_u = torch.tensor([35.7, 55.9, 58.2, 81.9, 56.3, 48.9,
                    33.9, 21.8, 48.4, 60.4, 68.4])

def model(t_u, w, b):
    return w * t_u + b

def loss_fn(t_p, t_c):
    squared_diffs = (t_p - t_c)**2
    return squared_diffs.mean()

params = torch.tensor([1.0, 0.0], requires_grad=True)

loss = loss_fn(model(t_u, params[0], params[1]), t_c)
loss.backward()

print(params.grad)

梯度计算逻辑解答 #

你的猜想完全正确，具体运行逻辑如下：

• PyTorch中对张量做索引得到的是原张量的视图（View），不是独立的新张量，和原张量共享底层数据存储，计算图会完整记录索引操作对应的梯度传播路径，不会因为取子张量就切断和原张量的关联。
• 前向传播过程中，params[0]（权重）、params[1]（偏置）参与计算时，计算图会分别记录损失对两个子张量的偏导数。
• 调用loss.backward()反向传播时，梯度会顺着计算图回传，因为两个子张量是params的视图，对应的梯度会自动写入params.grad的对应位置：params.grad[0]就是损失对权重的偏导，params.grad[1]就是损失对偏置的偏导，和单独定义两个requires_grad=True的w、b张量计算得到的梯度完全一致。

你可以自行验证：单独定义w = torch.tensor(1.0, requires_grad=True)、b = torch.tensor(0.0, requires_grad=True)，走完同样的前向、反向流程后打印w.grad和b.grad，得到的结果和params.grad的两个值完全相同。

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