autograd.function定义relu

下面是使用autograd.function定义ReLU函数的示例代码：

import torch

class ReLUFunction(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[input < 0] = 0
        return grad_input

x = torch.randn(5, 5, requires_grad=True)
y = ReLUFunction.apply(x)
z = y.sum()
z.backward()

print(x.grad)

在这个例子中，我们定义了一个名为ReLUFunction的autograd.Function。这个函数具有两个静态方法：forward和backward。在forward方法中，我们保存了输入张量，并对其进行了ReLU操作。在backward方法中，我们首先复制了梯度张量，然后将小于0的元素的梯度设置为0。最后，我们返回更新后的梯度张量。在使用这个自定义函数时，我们需要使用apply方法。在这个例子中，我们首先定义了一个需要求梯度的张量x，然后使用ReLUFunction.apply方法对它进行了ReLU操作，得到了y。最后，我们计算了y的和，并对z进行了反向传播，得到了x的梯度。