autograd.Function实现relu

下面是使用`autograd.Function`实现ReLU的示例代码：

import torch

class ReLUFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        # 保存输入，以便在反向传播时使用
        ctx.save_for_backward(input)
        # 对输入进行ReLU操作
        output = input.clamp(min=0)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        # 获取保存的输入
        input, = ctx.saved_tensors
        # 创建梯度tensor
        grad_input = grad_output.clone()
        # 将小于等于0的位置的梯度置为0
        grad_input[input <= 0] = 0
        return grad_input

# 定义输入
input = torch.randn(3, 3, requires_grad=True)
# 使用自定义函数进行ReLU操作
output = ReLUFunction.apply(input)
# 计算输出的梯度
output.sum().backward()

在这个例子中，我们定义了一个名为`ReLUFunction`的类，该类继承了`torch.autograd.Function`。该类包含两个静态方法：`forward`和`backward`，分别实现ReLU的前向传播和反向传播。

在`forward`方法中，我们保存了输入`input`，并对其进行了ReLU操作。在`backward`方法中，我们获取了保存的输入`input`和梯度`grad_output`，并计算了输入的梯度`grad_input`。具体来说，我们将小于等于0的位置的梯度置为0，因为在ReLU函数中，这些位置的梯度应该为0。

最后，我们使用自定义函数`ReLUFunction.apply`进行ReLU操作，并计算输出的梯度。