如何在PyTorch中直接访问原生函数的导数并支持反向传播？

PyTorch并没有将内置函数的反向计算逻辑暴露为独立的可直接调用函数，但你可以通过以下两种官方支持的方式，复用框架已实现的梯度逻辑，同时满足支持反向传播的需求，无需手动编写导数代码：

方法1：使用torch.autograd.functional.vjp（向量-雅克比乘积） #

vjp是PyTorch提供的直接调用目标函数反向传播逻辑的工具，它会返回一个可调用的梯度函数，该函数的计算过程完全复用内置函数的反向实现，且全程支持自动微分。

以sigmoid函数为例：

import torch

# 定义输入张量
x = torch.randn(3, requires_grad=True)

# 获取sigmoid函数的vjp：第一个返回值是sigmoid的输出，第二个是梯度计算函数
sigmoid_output, sigmoid_grad_fn = torch.autograd.functional.vjp(torch.sigmoid, x)

# 计算x的梯度：传入的张量对应损失函数对sigmoid输出的梯度（这里用全1张量等价于对输出求和后求导）
x_grad = sigmoid_grad_fn(torch.ones_like(sigmoid_output))[0]

# 验证结果与常规autograd一致
sigmoid_output.sum().backward()
assert torch.allclose(x_grad, x.grad)

这种方式完全不需要手动实现sigmoid的导数公式，直接复用框架内置的反向逻辑，且计算得到的x_grad本身也支持进一步的反向传播。

方法2：基于torch.autograd.Function封装 #

如果你需要一个更像“独立函数”的接口，可以基于PyTorch的Function类封装，内部复用内置函数的计算逻辑（包括反向）：

import torch

class SigmoidGradient(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        y = torch.sigmoid(x)
        ctx.save_for_backward(y)
        # 前向传播直接返回sigmoid对x的梯度
        return y * (1 - y)
    
    @staticmethod
    def backward(ctx, grad_output):
        y, = ctx.saved_tensors
        # 反向传播计算梯度的梯度（复用内置张量操作，自动支持微分）
        grad_x = grad_output * (y * (1 - y) * (1 - 2*y))
        return grad_x

# 使用封装后的函数
x = torch.randn(3, requires_grad=True)
grad = SigmoidGradient.apply(x)
# 对梯度求和后反向传播，验证二阶导数逻辑
grad.sum().backward()
print(x.grad)

注意：这种方式的前向计算虽然写了y*(1-y)，但本质是复用了sigmoid的输出结果，而反向过程的计算也依赖PyTorch内置的张量操作，无需手动推导高阶导数。

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