如何为PyTorch类Tensor类型实现自定义梯度计算？

问题1：原生PyTorch Tensor与自定义类互转时的梯度读写 #

直接访问tensor.grad是官方支持的合法梯度获取方式，不存在私有API访问的兼容性问题，注意几个实现细节即可：

• 读取原生Tensor梯度时，必须先做判空：只有设置requires_grad=True、且完成过反向传播计算的张量，grad属性才会有有效值，未触发梯度计算时该属性为None。
• 不要直接保存tensor.grad的引用，必须调用detach().clone()后再存入自定义存储，避免后续原地操作意外修改梯度值，破坏计算图。
• 从原生Tensor初始化自定义类的参考实现：

@classmethod
def from_torch(cls, t: torch.Tensor):
    obj = cls()
    # 读取并存储前向数值，必须detach避免绑定计算图
    obj._custom_data = t.detach().clone()
    # 读取并存储梯度
    if t.grad is not None:
        obj._custom_grad = t.grad.detach().clone()
    else:
        obj._custom_grad = None
    obj.requires_grad = t.requires_grad
    # 同步其他必要属性，比如dtype、device
    obj.dtype = t.dtype
    obj.device = t.device
    return obj

• 自定义类转换回原生Tensor时，不要直接给处于计算图中的张量赋值grad，正确流程是先构建数值张量、开启梯度标记，再同步梯度值：

def to_torch(self) -> torch.Tensor:
    # 构建输出张量，开启梯度标记
    out = self._custom_data.clone().requires_grad_(self.requires_grad)
    # 同步存储的梯度
    if self._custom_grad is not None:
        # 保证梯度和数值张量的device、dtype完全一致
        out.grad = self._custom_grad.to(device=self.device, dtype=self.dtype)
    return out

• 注意：不要对绑定了计算图的张量直接调用numpy()，必须先执行detach()，否则会触发梯度隔离报错。

问题2：为自定义张量类实现带反向逻辑的算子适配 #

自定义torch.autograd.Function的机制完全适配类Tensor类型的扩展场景，不需要修改PyTorch底层autograd逻辑，只要在__torch_function__分发逻辑中把对应算子路由到自定义Function即可，这也是官方示例采用的标准实现方案。

实现步骤 #

1. 编写对应算子的自定义autograd Function，分别实现自定义格式的前向计算、反向梯度逻辑：

class CustomAdd(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input1, input2, alpha=1):
        # input1、input2为自定义类实例，按自定义存储格式完成前向计算
        out_data = input1._custom_data + alpha * input2._custom_data
        # 保存反向传播需要的上下文信息
        ctx.alpha = alpha
        ctx.need_grad1 = input1.requires_grad
        ctx.need_grad2 = input2.requires_grad
        # 封装为自定义类实例返回
        out = MyCustomTensor()
        out._custom_data = out_data
        out.requires_grad = input1.requires_grad or input2.requires_grad
        return out

    @staticmethod
    def backward(ctx, grad_out):
        # grad_out为上游回传的梯度，类型为自定义类实例
        grad1 = None
        grad2 = None
        # 按自定义格式实现梯度计算逻辑，以add算子为例，梯度直接回传
        if ctx.need_grad1:
            grad1 = MyCustomTensor()
            grad1._custom_data = grad_out._custom_data.clone()
        if ctx.need_grad2:
            grad2 = MyCustomTensor()
            grad2._custom_data = ctx.alpha * grad_out._custom_data.clone()
        # 返回值数量和forward的入参数量严格对齐，不需要梯度的入参返回None
        return grad1, grad2, None

2. 在自定义类的__torch_function__方法中，将torch.add的调用路由到上述自定义Function的apply方法：

class MyCustomTensor:
    # 省略其他已实现的属性、方法
    @classmethod
    def __torch_function__(cls, func, types, args=(), kwargs=None):
        kwargs = kwargs if kwargs is not None else {}
        # 路由add算子
        if func is torch.add:
            return CustomAdd.apply(*args, **kwargs)
        # 其余算子按相同逻辑扩展
        raise NotImplementedError(f"算子 {func.__name__} 未完成自定义类型适配")

注意事项 #

• 自定义Function的forward、backward输入输出要保持类型统一，要么全程传递自定义类实例，要么在算子边界处完成自定义类和原生Tensor的转换，不要混合传参，否则会打断梯度流。
• 不要在__torch_function__中直接编写前向逻辑后手动拼接反向节点，路由到torch.autograd.Function是官方支持的、不会破坏计算图追踪的稳妥方案。
• 如果适配原地算子（如torch.add_），反向传播时需要额外处理输入张量的梯度引用，否则会出现梯度值错乱的问题。

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