torch.autograd.Function自定义反向求导

最新推荐文章于 2024-03-12 22:10:30 发布

harry_tea

最新推荐文章于 2024-03-12 22:10:30 发布

阅读量2k

点赞数

分类专栏： PyTorch 文章标签：计算机视觉 opencv python

本文链接：https://blog.csdn.net/weixin_41978699/article/details/123863775

版权

PyTorch 专栏收录该内容

50 篇文章 15 订阅

订阅专栏

Function自定义反向求导规则

Extending torch.autograd

Extending torch.autograd

在某些情况下我们的函数不可微（not differentiable），但是我们仍然需要对他求导时，就需要我们自定义求导方式，这里我们根据PyTorch官网给出的例子，来看一下torch.autograd.Function是如何运行的

官网给出的例子为LinearFunction，代码如下，这里我们假设输入为 $2\times3$ 的矩阵 $x$ ，权重也为 $2\times3$ 的矩阵 $w$ ，bias为 $\times 2$ 的矩阵 $b$ ，则

forward
$y = x*w^{T}+b$

backward

$\frac{\partial y}{\partial x} = w^{T} ~~~ \frac{\partial y}{\partial w} = x ~~~ \frac{\partial y}{\partial b} = 1$

from numpy import double
import torch
from torch.autograd import Function
# Inherit from Function
class LinearFunction(Function):

    # Note that both forward and backward are @staticmethods
    @staticmethod
    # bias is an optional argument
    def forward(ctx, input, weight, bias=None):
        ctx.save_for_backward(input, weight, bias)
        output = input.mm(weight.t())
        if bias is not None:
            output += bias.unsqueeze(0).expand_as(output)
        return output

    # This function has only a single output, so it gets only one gradient
    @staticmethod
    def backward(ctx, grad_output):
        # This is a pattern that is very convenient - at the top of backward
        # unpack saved_tensors and initialize all gradients w.r.t. inputs to
        # None. Thanks to the fact that additional trailing Nones are
        # ignored, the return statement is simple even when the function has
        # optional inputs.
        input, weight, bias = ctx.saved_tensors
        grad_input = grad_weight = grad_bias = None

        # These needs_input_grad checks are optional and there only to
        # improve efficiency. If you want to make your code simpler, you can
        # skip them. Returning gradients for inputs that don't require it is
        # not an error.
        if ctx.needs_input_grad[0]:
            grad_input = grad_output.mm(weight)
        if ctx.needs_input_grad[1]:
            grad_weight = grad_output.t().mm(input)
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum(0)

        return grad_input, grad_weight, grad_bias

input  = torch.tensor([[2.0, 1.5, 2.5], [1.0, 2.0, 3.0]], dtype=torch.double, requires_grad=True)
weight = torch.tensor([[3.0, 2.0, 3.5], [1.0, 2.0, 3.0]], dtype=torch.double, requires_grad=True)
bias   = torch.tensor([0.1, 0.2], dtype=torch.double, requires_grad=True)

# two ways to use linear operation
# first
output = LinearFunction.apply(input, weight, bias)
print(output)
# second
linear = LinearFunction.apply
output = linear(input, weight, bias)
print(output)


# 检查backward是否计算正确
from torch.autograd import gradcheck
# gradchek takes a tuple of tensor as input, check if your gradient
# evaluated with these tensors are close enough to numerical
# approximations and returns True if they all verify this condition.
test = gradcheck(LinearFunction.apply, (input, weight, bias), eps=1e-6, atol=1e-4)
print(test)  #　没问题的话输出True