pytorch_torch.autograd.Function

发表于 分类于 本文字数： 2.8k 阅读时长 ≈ 3 分钟

这是关于torch.autograd.Function

在 PyTorch 中，torch.autograd.Function 是一个基础类，用于定义自定义的autograd函数，使你能够实现任意的前向传播和反向传播操作。这对于实现自定义的操作和损失函数，或者对已有操作进行修改，都非常有用。

要使用 torch.autograd.Function，你需要创建一个继承自它的子类，并实现以下两个方法：forward 和 backward。

  1. forward 方法：
    这个方法定义了自定义函数的前向传播过程。它接收输入张量或其他变量作为参数，并返回计算结果。在 forward 方法中，你可以执行任意计算，包括创建新的张量和执行运算符。

  2. backward 方法：
    这个方法定义了自定义函数的反向传播过程。它接收关于输出的梯度（通常是一个梯度张量）作为参数，并计算相对于输入的梯度。在 backward 方法中，你需要计算输入变量的梯度，以便在整个计算图中进行梯度传播。

以下是一个简单的示例，演示如何使用 torch.autograd.Function 来实现一个自定义函数：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import torch

class MyFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        # 在 forward 方法中执行前向传播计算
        ctx.save_for_backward(input)
        output = input * 2
        return output

    @staticmethod
    def backward(ctx, grad_output):
        # 在 backward 方法中计算梯度
        input, = ctx.saved_tensors
        grad_input = grad_output * 2
        return grad_input

# 使用自定义函数
x = torch.tensor([1.0], requires_grad=True)
y = MyFunction.apply(x)
y.backward()

print("Input gradient:", x.grad)
print(y)
'''
Input gradient: tensor([2.])
tensor([2.], grad_fn=<MyFunctionBackward>)
'''

在这个示例中，MyFunction 继承自 torch.autograd.Function，并实现了 forward 和 backward 方法。你可以通过 MyFunction.apply() 来使用这个自定义函数。在后续的反向传播中，PyTorch 将会使用 backward 方法计算梯度。

这就是如何使用 torch.autograd.Function 来实现自定义函数，并在自定义的计算中使用 PyTorch 的自动微分。

  • @staticmethod 是 Python 中的一个装饰器（Decorator），用于将一个方法定义为静态方法。静态方法是指在类中定义的方法，不依赖于类的实例，因此可以直接通过类名调用，而不需要创建类的对象实例。

    在你提供的代码中，@staticmethod 装饰器用于将方法定义为静态方法。具体来说，它用于 SpecialSpmmFunction 类中的两个方法：forward 和 backward。

    1
    2
    3
    4
    5
    6
    7
    8
    class SpecialSpmmFunction(torch.autograd.Function):
        @staticmethod
        def forward(ctx, indices, values, shape, b):
            # ... implementation ...

        @staticmethod
        def backward(ctx, grad_output):
            # ... implementation ...

    通过将这两个方法定义为静态方法，你可以在不创建类的实例的情况下，直接通过类名调用这些方法。例如：

    1
    2
    3
    4
    5
    indices = ...
    values = ...
    shape = ...
    b = ...
    result = SpecialSpmmFunction.forward(indices, values, shape, b)

    这种方法非常适合在定义类的方法时，不需要访问实例属性或方法，或者在类的实例不存在的情况下执行一些操作。静态方法不会自动接收类的实例作为第一个参数（通常是 self），因此它们不依赖于类的状态。

  • 在上面的代码中，y = MyFunction.apply(x) 这一行代码是通过调用 MyFunction 类的 apply 方法来计算前向传播的结果 y。在这个特定的示例中，MyFunction 类的 forward 方法执行的操作是将输入张量 x 乘以 2，因此 y 的值将是 x 的两倍。

    这里，MyFunction.apply(x) 实际上是在前向传播中使用了自定义的操作，并返回计算得到的输出。因为我们定义了自定义函数 MyFunction 的 forward 方法，所以调用 .apply(x) 实际上就是调用了我们自己实现的操作。

    在更复杂的情况下，自定义函数可能会执行许多不同的操作，从而实现复杂的前向传播。apply 方法允许我们将输入传递给这些操作，并返回输出。通常情况下，PyTorch 的模块和函数也是这样工作的，只是在内部使用了更多的优化和组件。

    简而言之，y = MyFunction.apply(x) 将会调用自定义函数 MyFunction 的前向传播方法，执行该方法中的操作，并将操作的结果存储在 y 中。

  • 对于print(y)

    • 在上面的代码中，y = MyFunction.apply(x) 这一行代码是通过调用 MyFunction 类的 apply 方法来计算前向传播的结果 y。在这个特定的示例中，MyFunction 类的 forward 方法执行的操作是将输入张量 x 乘以 2，因此 y 的值将是 x 的两倍。

      这里，MyFunction.apply(x) 实际上是在前向传播中使用了自定义的操作，并返回计算得到的输出。因为我们定义了自定义函数 MyFunction 的 forward 方法，所以调用 .apply(x) 实际上就是调用了我们自己实现的操作。

      在更复杂的情况下，自定义函数可能会执行许多不同的操作，从而实现复杂的前向传播。apply 方法允许我们将输入传递给这些操作，并返回输出。通常情况下，PyTorch 的模块和函数也是这样工作的，只是在内部使用了更多的优化和组件。

      简而言之，y = MyFunction.apply(x) 将会调用自定义函数 MyFunction 的前向传播方法，执行该方法中的操作，并将操作的结果存储在 y 中。

  • 如果令c=y.backward(),print©输出的结果为None

  • 如果将y.backward()注释掉，print(“Input gradient:”, x.grad)为Input gradient:None