如何封装PyTorch函数并实现自动求导？

Question

我正在学习PyTorch教程，定义新的自动求导函数。我想要实现的自动求导函数是torch.nn.functional.max_pool1d的一个包装器。以下是我的代码：

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.autograd as tag

class SquareAndMaxPool1d(tag.Function):

    @staticmethod
    def forward(ctx, input, kernel_size, stride=None, padding=0, dilation=1, \
                return_indices=False, ceil_mode=False):
        ctx.save_for_backward( input )

        inputC = input.clone() #copy input
        inputC *= inputC

        output = F.max_pool1d(inputC, kernel_size, stride=stride, \
                              padding=padding, dilation=dilation, \
                              return_indices=return_indices, \
                              ceil_mode=ceil_mode)

        return output

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = get_max_pool1d_grad_somehow(grad_output)
        return 2.0*input*grad_input

我的问题是：如何得到包装函数的梯度？我知道可能有其他方法可以做到这一点，因为我提出的示例非常简单，但我想要做的适合这个框架，并且需要我实现一个autograd函数。

编辑：在查看此博客文章后，我决定尝试以下backward：

def backward(ctx, grad_output):
    input, output = ctx.saved_tensors
    grad_input = output.backward(grad_output)
    return 2.0*input*grad_input

将output加入保存的变量中。然后我运行以下代码:

x = np.random.randn(1,1,5)
xT = torch.from_numpy(x)
xT.requires_grad=True
f = SquareAndMaxPool1d.apply
s = torch.sum(f(xT,2))
s.backward()

我得到了总线错误：10。

假设xT是tensor([[[ 1.69533562, -0.21779421, 2.28693953, -0.86688095, -1.01033497]]], dtype=torch.float64)，那么我期望在调用s.backward()后找到xT.grad为tensor([[[ 3.39067124, -0. , 9.14775812, -0. , -2.02066994]]], dtype=torch.float64)（即2*x*grad_of_max_pool，其中grad_of_max_pool包含tensor([[[1., 0., 2., 0., 1.]]], dtype=torch.float64)）。

我已经找到了为什么出现“总线错误：10”的原因。看起来上面的代码导致我的backward在grad_input = output.backward(grad_output)处发生递归调用。所以我需要找到其他方法来获取max_pool1d的梯度。我知道如何在纯Python中实现这个，但结果会比包装库代码慢得多。

- Sean Lake

---

“get the gradient”是什么意思？实现？计算？ - Jatentaki

@Jatentaki 我的意思是，我相信PyTorch有一种方法可以计算所需的梯度，只要调用正确的函数即可。但我现在遇到了困难，无法确定该调用哪个函数。我刚刚在问题中添加了一个尝试的解决方案，但失败了。希望这能澄清事情。 - Sean Lake

Answer 1

你选择了一个不太幸运的例子。 torch.nn.functional.max_pool1d 不是 torch.autograd.Function 的实例，因为它是 PyTorch 内置的，在 C++ 代码中定义，并具有 autogenerated Python 绑定。我不确定是否可以通过其接口获取 backward 属性。

首先，如果您还没有注意到，您不需要编写任何自定义代码来反向传播此公式，因为幂运算和 max_pool1d 已经定义了它，所以它们的组合也被 autograd 所覆盖。假设您的目标是练习，我建议您手动执行它（而不回退到 max_pool1d 的 backward）。以下是一个示例

import torch
import torch.nn.functional as F
import torch.autograd as tag

class SquareAndMaxPool1d(tag.Function):
    @staticmethod
    def forward(ctx, input, kernel_size, **kwargs):
        # we're gonna need indices for backward. Currently SquareAnd...
        # never actually returns indices, I left it out for simplicity
        kwargs['return_indices'] = True

        input_sqr = input ** 2
        output, indices = F.max_pool1d(input_sqr, kernel_size, **kwargs)
        ctx.save_for_backward(input, indices)

        return output

    @staticmethod
    def backward(ctx, grad_output):
        input, indices = ctx.saved_tensors

        # first we need to reconstruct the gradient of `max_pool1d`
        # by putting all the output gradient elements (corresponding to
        # input elements which made it through the max_pool1d) in their
        # respective places, the rest has gradient of 0. We do it by
        # scattering it against a tensor of 0s
        grad_output_unpooled = torch.zeros_like(input)
        grad_output_unpooled.scatter_(2, indices, grad_output)

        # then incorporate the gradient of the "square" part of your
        # operator
        grad_input = 2. * input * grad_output_unpooled

        # the docs for backward
        # https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function.backward
        # say that "it should return as many tensors, as there were inputs
        # to forward()". It fails to mention that if an argument was not a
        # tensor, it should return None (I remember reading this somewhere,
        # but can't find it anymore). Anyway, we need to
        # return a (grad_input, None) tuple to avoid a complaint that two
        # outputs were expected
        return grad_input, None

我们可以使用数值梯度检查器来验证操作是否按预期工作。

f = SquareAndMaxPool1d.apply
xT = torch.randn(1, 1, 6, requires_grad=True, dtype=torch.float64)
tag.gradcheck(lambda t: f(t, 2), xT)

很抱歉这可能无法回答你如何获取max_pool1d的backward，但希望我的回答对你有所帮助。

- Jatentaki

---

谢谢！您说得对，这是一个练习，可以包装一个pytorch函数并在不知道该函数细节的情况下使autograd正常工作（或者是否可能）。我觉得我会采取另一种方法，将max_pool1d之前的部分放入一个autograd函数中，将其后的部分放入另一个函数中，但我发现您写的内容在其他方面也很有用和有帮助。 :) - Sean Lake

Answer 2

你在递归调用时遇到的问题实际上来自于output和默认情况下继承自torch.autograd.Function类的with no_grad行为。如果你在forward中检查output.grad_fn，它可能是None，而在backward中，它可能会链接到函数对象<SquareAndMaxPool1d...>，从而导致递归调用。如果你仍然有兴趣学习如何精确地执行你所要求的操作，这里有一个使用F.linear的示例：

import torch
import torch.nn.functional as F

class custom_Linear(nn.Linear):
    def forward(self, _input):
        return Custom_Linear_AGfn_getAround.apply(_input, self.weight, self.bias)

class Custom_Linear_AGfn_getAround(torch.autograd.Function):
    @staticmethod
    def forward(ctx, _input, _weight, _bias):
        print('Custom forward')
        with torch.enable_grad():
            detached_input = _input.detach()
            detached_input.requires_grad_(True)
            detached_weight = _weight.detach()
            detached_weight.requires_grad_(True)
            detached_bias = _bias.detach()
            detached_bias.requires_grad_(True)
            _tmp = F.linear(detached_input, detached_weight, detached_bias)
        ctx.saved_input = detached_input
        ctx.saved_param = detached_weight, detached_bias
        ctx.save_for_backward(_tmp)
        _output = _tmp.detach()
        return _output

    @staticmethod
    def backward(ctx, grad_out):
        print('Custom backward')
        _tmp, = ctx.saved_tensors
        _weight, _bias = ctx.saved_param
        detached_input = ctx.saved_input
        with torch.enable_grad():
            _tmp.backward(grad_out)
        return detached_input.grad, _weight.grad, _bias.grad

基本上，这只是关于为感兴趣的部分构建一个小的隔离图，而不干扰主图，并使用grad_fn和requires_grad来跟踪图形，在查看需要分离什么和隔离图需要什么时使用。

关于棘手的部分：

• 分离权重和偏置：您可以不使用，但是如果您通过save_for_backward传递_weight和_bias，则在backward内将_weight.grad，_bias.grad作为None传递，但是一旦外部_weight.grad，_bias.grad 将具有其正确值，或者您将它们通过属性（例如ctx.saved_param）传递，这种情况下，您将必须为backward的最后两个返回值手动放置None（返回detached_input.grad，None，None），否则您在之后检查权重和偏差梯度时将获得两倍的正确值。
• 如前所述，继承类torch.autograd.Function的backward和forward似乎默认具有with no_grad行为。因此，删除上述代码中的with torch.enable_grad()将导致_tmp.grad_fn为None（不明白为什么默认情况下，在forward中_tmp的grad_fn是None，而requires_grad是False尽管在detached_input需要梯度直到我遇到：https://github.com/pytorch/pytorch/issues/7698）
• 我相信但没有检查，如果您不将_output分离，则可能会获得双重grad_fn，因为当我没有使用with torch.enable_grad()并且不分离输出时（在前向传递中_tmp.grad_fn为None），它确实在backward中获取了<Custom_Linear_AGfn_getAround...> grad_fn（导致无限递归调用）。

- Romain Renard