我正在阅读有关定义新的 autograd 函数的 PyTorch 教程。我要实现的 autograd 函数是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()
我明白了Bus error: 10
。
说xT
is tensor([[[ 1.69533562, -0.21779421, 2.28693953, -0.86688095, -1.01033497]]], dtype=torch.float64)
,那么我希望在调用之后找到它(xT.grad
即 包含)。tensor([[[ 3.39067124, -0. , 9.14775812, -0. , -2.02066994]]], dtype=torch.float64)
s.backward()
2*x*grad_of_max_pool
grad_of_max_pool
tensor([[[1., 0., 2., 0., 1.]]], dtype=torch.float64)
我已经弄清楚为什么我得到一个Bus error: 10
. 上面的代码似乎导致了 my backward
at的递归调用grad_input = output.backward(grad_output)
。所以我需要找到其他方法来获得max_pool1d
. 我知道如何在纯 Python 中实现这一点,但结果会比我可以包装库代码要慢得多。