PyTorch中如何用grad函数实现exp_reducer对X指定列的偏导计算？

问题背景 #

需要计算函数exp_reducer对二维张量X的列元素组合后的偏导结果，预期输出与自定义derivative函数的输出一致。但直接调用torch.autograd.grad只能得到X每个元素的单独偏导数，需一套适用于高维复杂场景的通用解决方法。

原代码示例：

import numpy as np
import torch
from torch.autograd import grad 

X = torch.from_numpy(np.array([[2, 3],[4,5]],dtype = np.float32))
X.requires_grad = True

def derivative(X):
    res = 9* (X[:,0]**2 )* ( X[:,1]**2)
    return res

def exp_reducer(x):
    return x.pow(3).prod(axis = 1)

y = exp_reducer(X)

print(derivative(X)) # 预期输出
print(grad(y, X, grad_outputs=torch.ones_like(y))) # 当前仅能得到各元素单独偏导

逻辑分析 #

先明确exp_reducer的计算逻辑：对每个样本（行）的元素取3次方后相乘，即y_i = (X[i,0]^3) * (X[i,1]^3)。通过自动求导得到的梯度是X每个元素的单独偏导：

• 对X[i,0]的偏导：dy_i/dX[i,0] = 3*X[i,0]^2 * X[i,1]^3
• 对X[i,1]的偏导：dy_i/dX[i,1] = 3*X[i,0]^3 * X[i,1]^2

自定义derivative函数的输出9*X[i,0]^2*X[i,1]^2，可通过上述偏导推导得到：
9*X[i,0]^2*X[i,1]^2 = (dy_i/dX[i,0]) / X[i,1] * 3 = (dy_i/dX[i,1]) / X[i,0] *3

通用解决方案 #

核心思路是先获取全量梯度，再按自定义逻辑组合处理，步骤如下：

1. 调用torch.autograd.grad得到与X同形状的全量梯度张量。
2. 根据业务需求，对梯度张量进行索引、切片、数学运算等组合操作，得到预期结果。

针对本案例的实现代码：

# 获取与X同形状的全量梯度
full_grad = grad(y, X, grad_outputs=torch.ones_like(y))[0]

# 按自定义逻辑组合梯度，得到与derivative一致的结果
custom_grad = (full_grad[:,0] / X[:,1]) * 3
# 也可通过另一偏导验证：custom_grad = (full_grad[:,1] / X[:,0]) * 3

print(custom_grad)
# 输出：tensor([324., 3600.], grad_fn=<MulBackward0>)
# 与derivative(X)的输出完全一致

高维场景扩展建议 #

• 任意高维张量的梯度输出始终与输入形状一致，可直接利用PyTorch的张量运算（索引、广播、跨维度聚合等）完成组合。
• 若组合逻辑涉及复杂非线性变换，可直接基于梯度张量构建计算图，PyTorch会自动处理依赖关系。
• 多输出场景下，可将多个输出打包后调用grad，或分别求导后再进行组合计算。

内容的提问来源于stack exchange，提问作者G-09