自定义 Autograd 函数

TensorPlay 允许你通过定义自己的前向和反向逻辑来扩展其自动微分引擎。这对于实现带有自定义梯度的不可微操作或优化特定内核非常有用。

定义自定义函数

要创建自定义 autograd 函数，请继承 tpx.autograd.Function 并实现静态方法 forward 和 backward。

import tensorplay as tp
from tensorplay.autograd import Function

class MyExp(Function):
    @staticmethod
    def forward(ctx, i):
        result = i.exp()
        ctx.save_for_backward(result)
        return result

    @staticmethod
    def backward(ctx, grad_output):
        result, = ctx.saved_tensors
        return grad_output * result

# 使用方法
def my_exp(x):
    return MyExp.apply(x)

x = tp.randn(3, requires_grad=True)
y = my_exp(x)
y.sum().backward()

print(x.grad)

何时使用自定义函数？

1. 效率：如果你能比标准操作的组合更有效地计算组合梯度。
2. 稳定性：实现数值稳定版本的函数（例如 LogSumExp）。
3. 自定义硬件：对接具有自己微分逻辑的硬件。
4. 不可微操作：为 Step 函数或量化等操作提供“代理梯度”。

上下文对象 (ctx)

ctx 对象用于在正向传播和反向传播之间传递信息：

• ctx.save_for_backward(*tensors)：使用它来保存反向传播所需的张量。
• ctx.saved_tensors：在 backward 中访问保存的张量。
• ctx.mark_dirty(*tensors)：标记被就地修改的张量。
• ctx.mark_non_differentiable(*tensors)：标记不需要梯度的输出。