Как обернуть функции 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()

и получаю Bus error: 10.

Скажем, xT это 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.Похоже, что приведенный выше код приводит к рекурсивному вызову моего backward в grad_input = output.backward(grad_output).Поэтому мне нужно найти какой-то другой способ получить градиент max_pool1d.Я знаю, как реализовать это на чистом Python, но результат будет намного медленнее, чем если бы я мог обернуть код библиотеки.

Answer 1

Вы выбрали довольно неудачный пример.torch.nn.functional.max_pool1d не является экземпляром torch.autograd.Function, потому что это встроенный PyTorch, определенный в коде C ++ и с автоматически сгенерированным связыванием Python.Я не уверен, возможно ли получить свойство backward через его интерфейс.

Во-первых, если вы не заметили, вам не нужно писать какой-либо специальный код для обратного распространения этой формулы, потому чтокак для режима работы, так и для max_pool1d он уже определен, поэтому их состав также покрывается автоградой.Предполагая, что ваша цель - это упражнение, я бы посоветовал вам сделать это более вручную (без возврата к backward из max_pool1d).Ниже приведен пример

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)

Извините, еслиэто не решает ваш вопрос о том, как получить backward из max_pool1d, но, надеюсь, вы найдете мой ответ достаточно полезным.