Pytorch Custom Function用法

torch.autograd.Function

https://pytorch.org/docs/master/notes/extending.html

>>> class Exp(Function):
>>>
>>>     @staticmethod
>>>     def forward(ctx, i):
>>>         result = i.exp()
>>>         ctx.save_for_backward(result)
>>>         return result
>>>
>>>     @staticmethod
>>>     def backward(ctx, grad_output):
>>>         result, = ctx.saved_tensors
>>>         return grad_output * result

• 从pyotrch1.3开始，forward() backward() 都必须是@staticmethod.

backward(ctx, *grad_outputs)

• 默认第一个参数是ctx
• ctx后面跟的参数个数和forward() 的return个数相同（??）
• Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input.(??)
• gradient w.r.t (???)

Example ReLU pytorch defining new autograd functions

当需要对方程传入参数 non-Tensor arguments

class MulConstant(Function):
    @staticmethod
    def forward(ctx, tensor, constant):
        # ctx is a context object that can be used to stash information
        # for backward computation
        ctx.constant = constant
        return tensor * constant

    @staticmethod
    def backward(ctx, grad_output):
        # We return as many input gradients as there were arguments.
        # Gradients of non-Tensor arguments to forward must be None.
        return grad_output * ctx.constant, None