Как обернуть функции PyTorch и реализовать автоград?

Вы выбрали довольно неудачный пример. 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, но, надеюсь, вы найдете мой ответ достаточно полезным.

Проблемы, которые у вас были с рекурсивными вызовами, на самом деле исходят от output и того факта, что по умолчанию with no_grad является поведением по умолчанию, похоже, в объявлении класса, унаследованном от torch.autograd.Function. Если вы отметите output.grad_fn в forward, это, вероятно, будет 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 для отслеживания графов при поиске того, что нужно отсоединить и что необходимо для изолированного графа.

О сложных частях:

• отсоединение веса и смещения: вы можете обойтись без него, но ЛИБО вы затем передадите _weight и _bias через save_for_backward и получите _weight.grad, _bias.grad как None внутри backward, НО за пределами _weight.grad, _bias.grad будут иметь свои правильные значения, ИЛИ вы передадите их через атрибут как, скажем, ctx.saved_param, и в этом случае вам придется вручную указать None для двух последних возвращенных значений backward (return detached_input.grad, None, None), иначе вы получите в два раза правильное значение, когда впоследствии вы проверите градиент веса и смещения вне обратного.
• как было сказано в начале, backward и forward для унаследованного класса torch.autograd.Function, похоже, по умолчанию имеют поведение with no_grad. Таким образом, удаление with torch.enable_grad(): в приведенном выше коде приведет к тому, что _tmp.grad_fn будет None (Не могу понять, почему по умолчанию _tmp имел grad_fn до None и requires_grad до False в forward, несмотря на то, что для detached_input требовался градиент, пока я не наткнулся на: https://github.com/pytorch/pytorch/issues/7698)
• Я полагаю, но я не проверял, что вы можете получить двойной grad_fn для _output, если вы не отсоедините его, как когда у меня нет with torch.enable_grad(), и не отсоединяю выход, в результате чего _tmp.grad_fn в прямом направлении отсутствует, он получает <Custom_Linear_AGfn_getAround...> grad_fn в backward (и приводит к бесконечным рекурсивным вызовам).