pytorch 笔记: 扩展torch.autograd

Posted UQI-LIUWJ

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了pytorch 笔记: 扩展torch.autograd相关的知识,希望对你有一定的参考价值。

1 扩展torch.autograd

        向 autograd 添加操作需要为每个操作实现一个新的 Function 子类。

         回想一下,函数是 autograd 用来编码操作历史和计算梯度的东西。

2 何时使用

        通常,如果您想在模型中执行不可微分或依赖非 Pytorch 库中有的函数(例如 NumPy)的计算,但仍希望您的操作与其他操作链接并使用 autograd 引擎,可以通过扩展torch.autograd实现自定义函数 .

        在某些情况下,还可以使用自定义函数来提高性能和内存使用率:如果您使用扩展函数实现了前向和反向传递,则可以将它们包装在 Function 中以与 autograd 引擎交互。

3 何时不建议使用

        如果您已经可以根据 PyTorch 的内置操作编写函数,那么它的后向图(很可能)已经能够被 autograd 记录。 在这种情况下,您不需要自己实现反向传播功能。考虑使用一个普通的pytorch 函数即可。

4 如何使用

采取以下步骤:

1. 创建子类Function, 并实现 forward() 和 backward() 方法。

2. 在 ctx 参数上调用正确的方法。

3. 声明你的函数是否支持双反向传播(double backward)。

4. 使用 gradcheck 验证您的梯度(反向传播的实现)是否正确。

4.1 第一步

        在创建类 Function 之后,您需要定义 2 个方法:

  • forward() 是执行操作的代码(前向传播)

        它可以采用任意数量的参数,其中一些是optimal的(如果设定默认值的话)。

        这里接受各种 Python 对象。

        跟踪历史的张量参数(即 requires_grad=True 的tensor)将在调用之前转换为不跟踪历史的张量参数【但它们如何被使用将在计算图中注册】。请注意,此逻辑不会遍历列表/字典/任何其他数据结构,只会考虑作为调用的直接参数的张量。

        如果有多个输出,您可以返回单个张量输出或张量元组。

  • backward() 定义梯度公式。

      它将被赋予与和forward的输出一样多的张量参数,其中每个参数都代表对应输出的梯度。重要的是永远不要就地修改这些梯度(即不要有inplace操作)。

        它应该返回与forward输入一样多的张量,每个张量都包含对应输入的梯度。

        如果您的输入不需要梯度(needs_input_grad 是一个布尔元组,指示每个输入是否需要梯度计算),或者是非张量对象,您可以返回 None。

        此外,如果你有 forward() 的可选参数,你可以返回比输入更多的梯度,只要它们都是 None。

4.2 第二步

        需要正确使用 forward 的 ctx 相关函数,以确保新函数与 autograd 引擎一起正常工作。

  • save_for_backward() 可以 保存稍后在反向传播中需要使用的、前向的输入张量 。   

                任何东西,即非张量   和既不是输入也不是输出的张量,都应该直接存储在 ctx 上。

  • mark_dirty() 必须用于标记任何由 forward 函数就地修改的输入。
  • 使用 mark_non_differentiable() 来告诉autograd 引擎某一个输出是否不可微。默认情况下,所有可微分类型的输出张量都将设置为需要梯度。不可微分类型(即整数类型)的张量永远不会被标记为需要梯度。
  • set_materialize_grads() 可用于告诉 autograd 引擎在输出不依赖于输入的情况下优化梯度计算,方法是不物化给予后向函数的梯度张量。也就是说,如果设置为 False,python 中的 None 对象或 C++ 中的“未定义张量”(x.defined() 为 False 的张量 x)将不会在向后调用之前转换为填充零的张量,因此您的代码将需要处理这些对象,就好像它们是用零填充的张量一样。此设置的默认值为 True。  

4.3 第三步

        如果你的函数不支持两次反向传播double backward,你应该通过使用 once_differentiable() 向后修饰来明确声明它。 使用此装饰器,尝试通过您的函数执行双重反向传播double backward将产生错误。

4.4 第四步

        建议您使用 torch.autograd.gradcheck() 来检查您的后向函数是否正确计算前向梯度,方法是使用后向函数计算雅可比矩阵,并将值元素与使用有限差分数值计算的雅可比进行比较

5 示例

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

    @staticmethod
    # 注意这里forward和backward都是静态函数
    def forward(ctx, input, weight, bias=None):
        # bias 是一个可选变量,所以可以没有梯度

        ctx.save_for_backward(input, weight, bias)
        #这里就使用了step2中的save_for_backward
        #也就是保存稍后在反向传播中需要使用的、前向的输入或输张量

        output = input.mm(weight.t())

        if bias is not None:
            output += bias.unsqueeze(0).expand_as(output)
        return output
        #类似的前向传播输出定义

    
    @staticmethod
    def backward(ctx, grad_output):
        #由于这个方法只有一个输出(output),因而在backward中,只需要有一个输入即可(ctx不算的话)

        input, weight, bias = ctx.saved_tensors
        #这呼应的就是前面的ctx.save_for_backward

        grad_input = grad_weight = grad_bias = None

    
        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)
        '''
        这些 needs_input_grad 检查是可选的,只是为了提高效率。 
        如果你想让你的代码更简单,你可以跳过它们。 
        为不需要的输入返回梯度不会返回错误。 
        '''

        return grad_input, grad_weight, grad_bias

现在,为了更容易使用这些自定义操作,我们建议为它们的 apply 方法设置别名:

linear = LinearFunction.apply

这样之后,linear的效果就和我们正常的比如'loss=torch.nn.MSELoss'的loss差不多了 

5.1 用非tensor 参数化的方法

class MulConstant(Function):
    @staticmethod
    def forward(ctx, tensor, constant):
        ctx.constant = constant
        #非tensor的变量就不用save_for_backward了,直接存在ctx里面即可
        return tensor * constant

    @staticmethod
    def backward(ctx, grad_output):

        return grad_output * ctx.constant, None
        #非tensor 变量的梯度为0

6 检查效果

        您可能想检查您实现的反向传播方法是否实际计算了函数的导数。 可以通过与使用小的有限差分的数值近似进行比较:

from torch.autograd import gradcheck


input = (torch.randn(20,20,dtype=torch.double,requires_grad=True), 
         torch.randn(30,20,dtype=torch.double,requires_grad=True))
test = gradcheck(linear, input, eps=1e-6, atol=1e-4)
'''
gradcheck 将张量的元组作为输入,检查用这些张量评估的梯度是否足够接近数值近似值,

如果它们都验证了这个条件,则返回 True。 
'''
print(test)
#True

以上是关于pytorch 笔记: 扩展torch.autograd的主要内容,如果未能解决你的问题,请参考以下文章

pytorch 学习笔记之编写 C 扩展,又涨姿势了

PyTorch笔记 - MAE(Masked Autoencoders) PyTorch源码

PyTorch笔记 - MAE(Masked Autoencoders) PyTorch源码

PyTorch笔记 - MAE(Masked Autoencoders) PyTorch源码

PyTorch学习笔记:简介与基础知识

PyTorch笔记 - MAE: Masked Autoencoders Are Scalable Vision Learners