Pytorch的自定义拓展:torch.nn.Module和torch.autograd.Function

参考链接：pytorch的自定义拓展之（一）——torch.nn.Module和torch.autograd.Function_LoveMIss-Y的博客-CSDN博客_pytorch自定义backward前言：pytorch的灵活性体现在它可以任意拓展我们所需要的内容，前面讲过的自定义模型、自定义层、自定义激活函数、自定义损失函数都属于pytorch的拓展，这里有三个重要的概念需要事先明确。要实现自定义拓展，有两种方式，（1）方式一：通过继承torch.nn.Module类来实现拓展。这也是我们前面的例子中所用到的，它最大的特点是以下几点：包装torch普通函数和torch.nn...https://blog.csdn.net/qq_27825451/article/details/95189376pytorch的自定义拓展之（二）——torch.autograd.Function完成自定义层_LoveMIss-Y的博客-CSDN博客_pytorch 自定义function前言：前面的一篇文章中，已经很详细的说清楚了nn.Module、nn.functional、autograd.Function三者之间的联系和区别，虽然autograd.Function本质上是自定义函数的，但是由于神经网络、层、激活函数、损失函数本质上都是函数或者是多个函数的组合，所以使用autograd.Function依然可以达到定义层、激活函数、损失函数、甚至模型的目的，就像我们使...https://blog.csdn.net/qq_27825451/article/details/95312816定义torch.autograd.Function的子类，自己定义某些操作，且定义反向求导函数_tang-0203的博客-CSDN博客_saved_tensors大部分内容转载自： Pytorch入门学习（八）—–自定义层的实现（甚至不可导operation的backward写法） 哇,这个博客是对pytorch官方手册中-Extending PyTorch部分的的翻译总虽然pytorch可以自动求导，但是有时候一些操作是不可导的，这时候你需要自定义求导方式。也就是所谓的 “Extending torch.autograd”. 官网虽然...https://blog.csdn.net/tsq292978891/article/details/79364140

Pytorch入门学习（八）-----自定义层的实现（甚至不可导operation的backward写法）_Hungryof的博客-CSDN博客_pytorch自定义backwardpytorch自定义层，各种情况说明。https://blog.csdn.net/Hungryof/article/details/78346304

 由于各种知识的博客杂乱无章， 所以我就特意整理了一个比较完整的Pytorch的自定义拓展的博文。希望对所有人读了之后，都有所帮助。

前言：Pytorch 的灵活性体现在它可以任意扩展我们所需要的内容。以下内容都是属于Pytorch 的扩展。

• 自定义模型
• 自定义层
• 自定义激活函数
• 自定义损失函数

要实现自定义拓展，有两种方式。

方式一：通过继承torch.nn.Module 类来实现扩展。特点有：

• 包装torch普通函数和torch.nn.functional专用于神经网络的函数；（torch.nn.functional是专门为神经网络所定义的函数集合）
• 只需要重新实现__init__和forward函数，求导的函数是不需要设置的，会自动按照求导规则求导(Module类里面是没有定义backward这个函数的）
• 可以保存参数和状态信息；

方式二：通过继承torch.nn.Function类来实现扩展。特点有：

• 在有些操作通过组合Pytorch中已有的层或者是已有的方法实现不了的时候，比如你要实现一个新的方法，这个新的方法需要forward和backward一起写，然后自己写对中间变量的操作。
• 需要重新实现__init__和forward函数，以及backward函数，需要自己定义求导规则；
• 不可以保存参数和状态信息

总结：当不使用自动求导机制，需要自定义求导规则的时候，就应该拓展torch.autograd.Function类。 否则就是用torch.nn.Module类，后者更简单更常用。

一、为什么要使用torch.nn.Function类

        Pytorch 中有自动求导机制，但是这都是针对torch里面所定义的函数，从上面已经知道了torch.nn.functional是专门为神经网络所定义的函数集合。如果一些操作torch.nn.functional没有提供，而torch里面也没有提供，那肿么办？

        当然我们可以使用一些基本的Pytorch 函数来进行组装，另外我们也可以使用numpy或scipy三方库中的方法实现。这个时候由于Pytorch 不再提供自动求导机制，就要自己定义实现前向传播和反向传播的计算过程了。

        另外，虽然Pytorch 可以自动求导，但是有时候一些操作是不可导的，这时候你需要自定义求导方式。也就是所谓的 “Extending torch.autograd”。

1.1、torch.autograd.Function类的定义

class Function(with_metaclass(FunctionMeta, _C._FunctionBase, _ContextMethodMixin, _HookMixin)):
 
    __call__ = _C._FunctionBase._do_forward
    is_traceable = False
 
    @staticmethod
    def forward(ctx, *args, **kwargs):
 
    @staticmethod
    def backward(ctx, *grad_outputs):

里面的小例子：

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

#Use it by calling the apply method:
output = Exp.apply(input)

其实就是实现forward和backward两个函数。注意这里和Module类最明显的区别是它多了一个backward方法，这也是他俩最本质的区别：

1. torch.autograd.Function类实际上是某一个操作函数的父类，一个操作函数必须具备两个基本的过程，即前向的运算过程和反向的求导过程，
2. torch.nn.Module类实际上是对torch.xxxx以及torch.nn.functional.xxxx这些函数的包装组合，而torch.xxxx和torch.nn.functional.xxxx都是实现了torch.autograd.Function类的两个基本功能（前向运算和反向传播），如果是我们需要的某一个功能torch.xxxx和torch.nn.functional里面都没有，也不能通过组合得到，这就需要定义新的操作函数，这个函数就需要继承自autograd.Function类，重写前向运算和反向传播。这里的torch.autograd.Function是一个最基本的，最底层的一个类。所有需要forward和backward都需要继承这个类（注意体会这段话）
3. 很显然，nn.Module更加高层，而autograd.Function更加底层，其实从名字中也能看出二者的区别，Module是针对模块的，即神经网络中的层、激活层、损失函数、网络模型等等，而Function是针对函数的，针对的是一些需要自己定义的函数而言的。如果某一个函数myFunc继承自Function类，实现了这个类的forward和backward方法，那么我依然可以用nn.Module对这个自定义的的函数myFunc进行包装组合，因为此时myFunc跟torch.xxxx和torch.nn.functional.xxxx里面的函数已经具备了等同的地位了。（注意体会这段话），可以这么说，Module不仅包括了Function，还包括了对应的参数，以及其他函数与变量，这是Function所不具备的。
4. 那为什么Function类也可以定义一个神经网络

from torch.autograd import Variable
class MyReLU(torch.autograd.Function):
    @staticmethod
    def forward(self, input_):
        # 在forward中，需要定义MyReLU这个运算的forward计算过程；
        # 同时可以保存任何在后向传播中需要使用的变量值
        self.save_for_backward(input_) # 将输入保存起来，在backward时使用
        output = input_.clamp(min=0)      # relu就是截断负数，让所有负数等于0
        return output
    @staticmethod
    def backward(self, grad_output):
        # 根据BP算法的推导(链式法则)，dloss / dx = (dloss / doutput) * (doutput / dx)
        # dloss / doutput就是输入的参数grad_output
        # 因此只需求relu的导数，在乘以grad_output
        input_, = self.saved_tensors
        grad_input = grad_output.clone()
        grad_input[input_ < 0] = 0 # 上诉计算的结果就是左式。即ReLU在反向
        #传播中可以看做一个通道选择函数，所有未达
        #到阈值(激活值<0)的单元的梯度都为0
        return grad_input

# 验证Variable与Function的关系from torch.autograd import Variable
input_=Variable(torch.randn(1))
relu=MyReLU.apply
output_=relu(input_)# 这个relu对象，就是output_.creator，即这个relu对象将

        很显然我们使用Function类自定义了一个神经网络模型，其实这么理解就好了，那就是：神经网络本质上来说就是一个较复杂的函数，它是由很多的函数运算组合起来的一个复杂函数，所以这里的MyReLU本质上来说还是一个torch的函数，而且我们可以看见，这个模型MyReLU是没有参数信息和状态信息保留的。

        所以如果我们现在使用autograd.Function类来自定义一个模型、一个层、一个激活函数、一个损失函数，就更加好理解了，实际上本质上来说都是一个函数，只分这个函数是简单还是复杂。

注意：这里必须使用MyReLU.apply。以前的是MyReLU()这不过随着Pytorch的升级而改变了

1.2、小结

有了上面这几点认识，我们可以概括性的得出这几样结论

（1）torch.nn.Module和torch.autograd.Function都是为Pytorch提供自定义拓展的途径；

（2）二者可以实现极度类似的功能，但二者所处的位置却完全不一样，二者的本质完全不一样；
 

二、自定义实现继承autograd.Function类

        由于这个类（torch.autograd.Function）确实是比较底层，正在使用的时候经常遇见我找不到的原因，所以本文只列举较为简单的情况，即不使用torch之外的三方库（numpy、scipy等，由于numpy和scipy函数是不支持backward的，所以在使用的时候涉及到ndarray与tensor之间的转换，常常出错），另外也暂时不涉及向量对向量的求导，仅仅涉及标量对标量和标量对向量求导，

2.1、标量对标量求导

举例子：

z=sqrt(x)+1/x+2*power(y,2)

z是关于x,y的一个二元函数它的导数是

z'(x)=1/(2*sqrt(x))-1/power(x,2)

z'(y)=4*y

import torch
# 定义一个继承了Function类的子类，实现y=f(x)的正向运算以及反向求导
class basicFunc(torch.autograd.Function):
    '''
    forward和backward可以定义成静态方法，向定义中那样，也可以定义成实例方法
    '''
    # 前向运算
    @staticmethod
    def forward(ctx, input_x, input_y):
        '''
        self.save_for_backward(input_x,input_y) ,这个函数是定义在Function的父类_ContextMethodMixin中
             它是将函数的输入参数保存起来以便后面在求导时候再使用，起前向反向传播中协调作用
        '''
        ctx.save_for_backward(input_x, input_y)
        output = torch.sqrt(input_x) + torch.reciprocal(input_x) + 2 * torch.pow(input_y, 2)
        return output
    @staticmethod
    def backward(ctx, grad_output):
        input_x, input_y = ctx.saved_tensors  # 获取前面保存的参数,也可以使用self.saved_variables
        if ctx.needs_input_grad[0]:
            grad_x = grad_output * (torch.reciprocal(2 * torch.sqrt(input_x)) - torch.reciprocal(torch.pow(input_x, 2)))
        if ctx.needs_input_grad[1]:
            grad_y = grad_output * (4 * input_y)
        return grad_x, grad_y  # 需要注意的是，反向传播得到的结果需要与输入的参数相匹配

# 由于sqrt_and_inverse是一个类，我们为了让它看起来更像是一个pytorch函数，需要包装一下
def sqrt_and_inverse_func(input_x, input_y):
    return basicFunc.apply(input_x, input_y)  # 这里是对象调用的含义，因为function中实现了__call__

x = torch.tensor(3.0, requires_grad=True)  # 标量
y = torch.tensor(2.0, requires_grad=True)
print('开始前向传播')
z = sqrt_and_inverse_func(x, y)
print('开始反向传播')
z.backward()  # 这里是标量对标量求导
print(x.grad)
print(y.grad)
'''运行结果为：
开始前向传播
开始反向传播
tensor(0.1776)
tensor(8.)
'''

2.2、标量向量求导

z=sum(sqrt(x*x-1)

这个时候x是一个向量，x=[x1,x2,x3]

则:z'(x)=x/sqrt(x*x-1)

import torch
class basicFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input_x):  # input_x是一个tensor，不再是一个标量
        ctx.save_for_backward(input_x)
        output = torch.sum(torch.sqrt(torch.pow(input_x, 2) - 1))  # 函数z
        return output
    @staticmethod
    def backward(ctx, grad_output):
        input_x, = ctx.saved_tensors  # 获取前面保存的参数,也可以使用self.saved_variables  #input_x前面的逗号是不能丢的
        if ctx.needs_input_grad[0]:
            grad_x = grad_output * (torch.div(input_x, torch.sqrt(torch.pow(input_x, 2) - 1)))
        return grad_x

def sqrt_and_inverse_func(input_x):
    return basicFunc.apply(input_x)  # 对象调用

x = torch.tensor([2.0, 3.0, 4.0], requires_grad=True)  # tensor
print('开始前向传播')
z = sqrt_and_inverse_func(x)
print('开始反向传播')
z.backward()
print(x.grad)
'''运行结果为：
开始前向传播
开始反向传播
tensor([1.1547, 1.0607, 1.0328])
'''

2.3、使用autograd.Function进行拓展的一般模板

import torch.nn
from torch.autograd import Function
class basicFunc(Function):
    @staticmethod
    def forward(self, inputs, parameters):
        self.saved_for_backward(inputs, parameters) # 可能在backward 中需要
        # output = [对inputs和 =parameters进行的操作，其实就是前向运算的函数表达式]
        return output
    @staticmethod
    def backward(self, grad_output):
        inputs, parameters = self.saved_tensors  # 或者是self.saved_variables
        # grad_inputs = [求函数forward(input)关于 parameters 的导数，其实就是反向运算的导数表达式] * grad_output
        return grad_input

# 包装自定义的myFunc有几种方法，通过方法包装，通过一个类包装都可以

# 这里就展示使用一个方法包装
# 这样使得看起来更加自然，因为Function的作用就是实现一个自定义方法的
def myFunc(inputs):
    return basicFunc.apply(inputs)  # 一定要是对象调用

# 这里就展示使用一个类包装
class myFunc(torch.nn.Module):
    def __init__(self):
        super(myFunc, self).__init__()

    def forward(self, x):
        basicFunc.apply(x)

'''注意事项：
需要注意的是，这里一定要使用对象调用，否则虽然也能够求出倒数结果，但实际上跟我自己定义backward函数就没啥关系了
如果使用 return myFunc.apply.forward(inputs)
这是不行的，虽然结果正确，后面会分析
'''

 2.4、自定义类继承自Function类的两个注意点

（1）注意点一：关于“对象调用”

包装函数里面一定要使用

basicFunc.apply(inputs) 

即对象调用，而不能使用，basicFunc.apply.forward(inputs),为什么？看下面的例子，依然以第上面的2.2例子而言，将backward改为如下：

def backward(self, grad_output):  
    print("---------------------------------------------") 
    print(f"grad_output is : {grad_output}")    
                          
    input_x,=self.saved_variables  #input_x前面的逗号是不能丢的
    if ctx.needs_input_grad[0]:
        grad_x = grad_output * (torch.div(input_x, torch.sqrt(torch.pow(input_x, 2) - 1)))
    return grad_x

如果包装函数如下：

def sqrt_and_inverse_func(input_x):
    return sqrt_and_inverse.apply(input_x)  #对象调用
'''运行结果为：
开始前向传播
开始反向传播
---------------------------------------------
grad_output is : 1.0
tensor([1.1547, 1.0607, 1.0328])
'''

 从上面可见我自己定义的backward的的确确是调用了的，如果我改为下面：

def sqrt_and_inverse_func(input_x):
    return sqrt_and_inverse.apply.forward(input_x) # 不是对象调用了
 
'''
开始前向传播
开始反向传播
tensor([1.1547, 1.0607, 1.0328])
'''

        我们发现自己定义的backward函数根本没有使用，虽然结果是一样的，为什么会这样子？

        其实第二种方法中，仅仅是调用了forward函数，而这个forward函数里面又定义了几个普通torch函数组合而成，所以实际上求导是直接对forward里面的那个表达式求导，但是由于我上面本来就是使用的简单torch函数，他们本来就是可以求导的，所以依然会得到相同的结果，而并不是通过自己定义的backward来实现的。所以上面的包装一定要通过“对象调用”来实现。

（2）注意点二：关于backward函数里面的grad_output参数

        通过上面的注意点一，在上面的两个例子中，例子2.1、2.2中我们得到的grad_output参数是1，这是为什么？要把这个问题交代清楚，需要一步一步来看，前面的一片文章提到过如果是向量对向量求导，需要给y.backward函数传递一个和被求导向量维度一样的tensor作为参数，backward的定义如下：

backward(gradient=None, retain_graph=None, create_graph=False)

而在我们自己定义的函数（继承自Function的类）里面的backward函数的定义如下：

def backward(ctx, grad_output):

        其实这里的grad_output实际上就是上面的gradient参数，本文的例子中，由于是标量对标量、标量对向量求导，所以没有传递这个grad_output参数，默认值就是1，这也就是上面为什么是1的原因，当然我可以给这个backward传递一个新的参数，如下：

gradient=torch.tensor(2.5)
z.backward(gradient)   # 这里是标量对标量求导，注意这个参数一定要是一个tensor才行
'''运行结果为：
开始前向传播
开始反向传播
---------------------------------------------
grad_output is : 2.5   # 这个时候grad_output的值就是我传递进去的2.5了
tensor(0.4439)         # 原来的 0.1776*2.5=0.4439
tensor(20.)            # 原来的 8.0*2.5=20.0
'''

小结：自定义函数backward中的grad_output实际上就是通过backward传递进去的参数gradient，这个参数必须是一个tensor类型，当是标量求导的时候，它是一个标量值，当是向量求导的时候，它是一个和求导向量同维度的向量（python的广播）。

那为什么是这样子呢？我似乎没有显示得调用自定义类的backward函数啊，我们来简单分析一下：

print('开始前向传播')
z=sqrt_and_inverse_func(x,y)  
print(z)
print(z.grad_fn)  
'''运行结果为：
开始前向传播
tensor(10.0654, grad_fn=<sqrt_and_inverse>)
<__main__.sqrt_and_inverse object at 0x000002AD04C75848>
'''

        我们发现这里的z是通过我们自己所定义的函数来创建出来的，pytorch中每一个tensor都有一个 grad_fn 属性，表示是谁创造了它，从这里可以看出，z 是由sqrt_and_inverse 创造出来的，所以调用z.backward()就是调用了sqrt_and_inverse.backward()，这也就是为什么编辑器中，将鼠标悬停在z.backward()上面却显示它的定义是sqrt_and_inverse.backward()的原因了。

补充：关于tensor的grad_fn属性：

每个tensor都有一个“.grad_fn”属性，这个属性表示的含义是谁创造了这个“Tensor”，如果是用户自己创造的，grad_fn属性就是None,否则就指向创造这个tensor的操作，如下：

import torch
x = torch.tensor(torch.ones(2,2),requires_grad=True)
y=x+2
print(x.grad_fn)  # 返回 None
print(y.grad_fn)  # 返回 <AddBackward object at 0x0000> 表示是由Add加法创造得到的Y

也可以参考我在1.1中的类MyReLU

根据BP算法的推导(链式法则)，dloss / dx = (dloss / doutput) * (doutput / dx)
dloss / doutput就是输入的参数grad_output
 因此只需求relu的导数，在乘以grad_output

三、具体的示例

        对于复杂的层或者是网络，使用autograd.Function几乎是不可行的，因为我们需要重新定义反向求导规则即backward函数，而复杂层或者网络没办法写出每一个参数的导函数，或者是即便写出来也是异常复杂（因为链式求导法则再加上一些非线性函数的关系）所以一般不推荐使用autograd.Function去定义层，更不要去定义模型，但是一般定义一个较简单的函数还是可以的。

3.1 重写Function类的静态方法

import torch
from torch.autograd import Function
 
# 类需要继承Function类，此处forward和backward都是静态方法
class MultiplyAdd(Function):  
                                                             
    @staticmethod                                  
    def forward(ctx, w, x, b):                 
        ctx.save_for_backward(w,x)    #保存参数,这跟前一篇的self.save_for_backward()是一样的
        output = w * x + b
        return output                        
         
    @staticmethod                                 
    def backward(ctx, grad_output):    #获取保存的参数,这跟前一篇的self.saved_variables()是一样的
        w,x = ctx.saved_variables  
        print("=======================================")             
        grad_w = grad_output * x
        grad_x = grad_output * w
        grad_b = grad_output * 1
        return grad_w, grad_x, grad_b  # backward输入参数和forward输出参数必须一一对应
 
x = torch.ones(1,requires_grad=True)  # x 是1，所以grad_w=1
w = torch.rand(1,requires_grad=True)  # w 是随机的，所以grad_x=随机的一个数
b = torch.rand(1,requires_grad=True)  # grad_b 恒等于1
 
print('开始前向传播')
z=MultiplyAdd.apply(w, x, b)   # forward,这里的前向传播是不一样的，这里没有使用函数去包装自定义的类，而是直接使用apply方法
print('开始反向传播')
z.backward()                   # backward
 
print(x.grad, w.grad, b.grad)
'''运行结果为：
开始前向传播
开始反向传播
=======================================
tensor([0.1784]) tensor([1.]) tensor([1.])
'''

注意：上面最大的不同除了使用的是静态方法以外，最大的不同在于，我没有使用一个函数去包装我的自定义类，而是直接使用了  z=MultiplyAdd.apply(w, x, b)  去完成前向运算过程，

这个apply方法是定义在Function类的父类_FunctionBase中定义的一个方法，但是这个方法到底是怎么实现的还不得而知。

四、具体分析

4.1、一些简单的参数定义

属性（成员变量）
saved_tensors: 传给forward()的参数，在backward()中会用到。
needs_input_grad:长度为 :attr:num_inputs的bool元组，表示输出是否需要梯度。可以用于优化反向过程的缓存。
num_inputs: 传给函数 :func:forward的参数的数量。
num_outputs: 函数 :func:forward返回的值的数目。
requires_grad: 布尔值，表示函数 :func:backward 是否永远不会被调用。

成员函数
forward()
forward可以有任意多个输入、任意多个输出，但是输入和输出必须是Variable。(官方给的例子中有只传入tensor作为参数的例子)
backward()
backward的输入和输出的个数就是forward函数的输出和输入的个数。其中，backward输入表示关于forward输出的梯度(计算图中上一节点的梯度)，backward的输出表示关于forward的输入的梯度。在输入不需要梯度时（通过查看needs_input_grad参数）或者不可导时，可以返回None。

ctx is a context object that can be used to stash information for backward computation（ctx是一个上下文对象，可用于存储向后计算的信息）

这里自己定义一个线性函数(传入参数是Variable)
 

y = x*w +b # 自己定义的LinearFunction
z = f(y)类似于z=loss(y)
下面的grad_output = dz/dy
根据复合函数求导法则:
1. dz/dx =  dz/dy * dy/dx = grad_output*dy/dx = grad_output*w
2. dz/dw =  dz/dy * dy/dw = grad_output*dy/dw = grad_output*x
3. dz/db = dz/dy * dy/db = grad_output*1

import torch
from torch.autograd import Function
from torch.autograd import Variable
class LinearFunction(Function):
    # 创建torch.autograd.Function类的一个子类
    # 必须是staticmethod
    @staticmethod
    # 第一个是ctx，第二个是input，其他是可选参数。
    # ctx在这里类似self，ctx的属性可以在backward中调用。
    # 自己定义的Function中的forward()方法，所有的Variable参数将会转成tensor！因此这里的input也是tensor．在传入forward前，autograd engine会自动将Variable unpack成Tensor。
    def forward(ctx, input, weight, bias=None):
        print(type(input))
        ctx.save_for_backward(input, weight, bias) # 将Tensor转变为Variable保存到ctx中
        output = input.mm(weight.t())  # torch.t()方法，对2D tensor进行转置
        if bias is not None:
            output += bias.unsqueeze(0).expand_as(output)
            # expand_as(tensor)等价于expand(tensor.size()), 将原tensor按照新的size进行扩展
        return output

    @staticmethod
    def backward(ctx, grad_output):
        # grad_output为反向传播上一级计算得到的梯度值
        input, weight, bias = ctx.saved_variables
        grad_input = grad_weight = grad_bias = None
        # 分别代表输入,权值,偏置三者的梯度
        # 判断三者对应的Variable是否需要进行反向求导计算梯度
        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).squeeze(0)

        return grad_input, grad_weight, grad_bias

#把新操作封装在一个类中
class linearModule(torch.nn.Module):
    def __init__(self):
        super(linearModule, self).__init__()

    def forward(self, inputs,weights, bias=None):
        return LinearFunction.apply(inputs,weights, bias)  # 调用forward()

#建议把新操作封装在一个函数中
def linear(inputs, weights, bias=None):
    # First braces create a Function object. Any arguments given here
    # will be passed to __init__. Second braces will invoke the __call__
    # operator, that will then use forward() to compute the result and
    # return it.
    return LinearFunction.apply(inputs, weights, bias)#调用forward()

# 或者使用apply方法对自己定义的方法取个别名
linear = LinearFunction.apply

#检查实现的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.

inputs = torch.randn((4,4), requires_grad=True, dtype=torch.float64)
weights = torch.randn((4,4), requires_grad=True, dtype=torch.float64)
test = gradcheck(linearModule(), (inputs,weights), eps=1e-6, atol=1e-4)
print(test)  #　没问题的话输出True

这里定义一个乘以常数的操作(输入参数是Tensor)

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.
        # constant
        return grad_output * ctx.constant, None # 这里并没有涉及到Variable

4.2、用自己定义的Function来创建Module

扩展module就很简单，需要重载 nn.Module中的init和forward

import torch.nn as nn
class Linear(nn.Module):
    def __init__(self, input_features, output_features, bias=True):
        super(Linear, self).__init__()
        self.input_features = input_features
        self.output_features = output_features
        # nn.Parameter is a special kind of Variable, that will get
        # automatically registered as Module's parameter once it's assigned
        # 这个很重要！ Parameters是默认需要梯度的！
        self.weight = nn.Parameter(torch.Tensor(output_features, input_features))
        if bias:
            self.bias = nn.Parameter(torch.Tensor(output_features))
        else:
            # You should always register all possible parameters, but the
            # optional ones can be None if you want.
            self.register_parameter('bias', None)
        # Not a very smart way to initialize weights
        self.weight.data.uniform_(-0.1, 0.1)
        if bias is not None:
            self.bias.data.uniform_(-0.1, 0.1)
    def forward(self, input):
        # See the autograd section for explanation of what happens here.
        return LinearFunction.apply(input, self.weight, self.bias)
        # 或者　return LinearFunction()(input, self.weight, self.bias)

4.3、forward的具体说明

1. 虽然说一个网络的输入是Variable形式，那么每个网络层的输出也是Variable形式。但是，当自定义autograd时，在forward中，所有的Variable参数将会转成tensor！因此在forward实际操作的对象是tensor。在传入forward前，autograd engine会自动将Variable unpack成Tensor。因此这里的input也是tensor.在forward中可以进行任意操作。
2. ctx是context，ctx.save_for_backward会将他们转换为Variable形式。也就是说, backward只对Variable进行处理.
3. save_for_backward只能传入Variable或是Tensor的变量，如果是其他类型的，可以用ctx.constant = constant ，使其在backward中可以用。例如,上面的ctx.constant = constant,这里constant为常数，不能直接作为ctx.save_for_backward的参数.

4.4、backward的具体说明

        1.grad_output是variable

        看到grad_output时候，会发现它是一个Variable，至于requires_grad是否为True，取决于你在外面调用.backward或是.grad时候的那个Variable是不是需要grad的。如果那个Variable是需要grad的，那么我们这里反向的grad_ouput也是requires_grad为True，那么我们甚至可以计算二阶梯度

        2.backward中我能一开始就.data拿出数据进行操作吗？

        虽然自定义操作，但是原则上在backward中我们只能进行Variable的操作, 这显然就要求我们在backward中的操作都是可自动求导的。所以如果我们的涉及到不可导的操作，那么我们就不能在backward函数中创建一个正确的图。
        3.自动求导是根据每个op的backward创建的graph来进行的
        大家的求导想法应该是：forward记录input如何被操作，然后backward就会自动反向，根据forward创建的图进行！然而，当你print(type(input))时你竟然发现类型是Tensor，根本不是Variable！那怎么记录graph？然而真实情况竟然是自动求导竟然是在backward的操作中创建的图！
        这也就是为什么我们需要在backward中用全部用variable来操作，而forward就没必要，forward只需要用tensor操作就可以。
        4.non-differential操作的backward怎么写？
      一般直接backward中第一句就是grad_output = grad_output.data,这样我们就无法进行创建正确的graph了。

from torch.autograd.function import once_differentiable
@staticmethod
@once_differentiable
def backward(ctx, grad_output):
    print(type(grad_output)) # 此时你会惊奇的发现，竟然是Tensor了！
    # 对grad_output 进行系列操作，得到grad_output_changed
    grad_input = grad_output_changed
    return grad_input

因为我们在backward中已经是直接拿出data进行操作的了，所以我们直接得到Tensor类型返回就行！

可以看non-differential的例子。具体没有仔细研究

torch.autograd.function.FunctionCtx.set_materialize_grads — PyTorch 1.11.0 documentationhttps://pytorch.org/docs/stable/generated/torch.autograd.function.FunctionCtx.set_materialize_grads.html

from torch.autograd import Function
from torch.autograd.function import once_differentiable
import torch
class SimpleFunc(Function):
    @staticmethod
    def forward(ctx, x):
        return x.clone(), x.clone()
    @staticmethod
    @once_differentiable
    def backward(ctx, g1, g2):
        return g1 + g2  # No check for None necessary
# We modify SimpleFunc to handle non-materialized grad outputs
class Func(Function):
    @staticmethod
    def forward(ctx, x):
        ctx.set_materialize_grads(False)
        ctx.save_for_backward(x)
        return x.clone(), x.clone()
    @staticmethod
    @once_differentiable
    def backward(ctx, g1, g2):
        x, = ctx.saved_tensors
        grad_input = torch.zeros_like(x)
        if g1 is not None:  # We must check for None now
            grad_input += g1
        if g2 is not None:
            grad_input += g2
        return grad_input
a = torch.tensor(1., requires_grad=True)
print(SimpleFunc.apply(a))
b, _ = Func.apply(a)  # induces g2 to be undefined

4.5、两个小示例

# -*- coding: utf-8 -*-
import torch
from torch.autograd import Variable
class MyReLU(torch.autograd.Function):
    """
    We can implement our own custom autograd Functions by subclassing
    torch.autograd.Function and implementing the forward and backward passes
    which operate on Tensors.
    """
    @staticmethod
    def forward(ctx, input):
        """
        In the forward pass we receive a Tensor containing the input and return
        a Tensor containing the output. ctx is a context object that can be used
        to stash information for backward computation. You can cache arbitrary
        objects for use in the backward pass using the ctx.save_for_backward method.
        """
        ctx.save_for_backward(input)
        return input.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_output):
        """
        In the backward pass we receive a Tensor containing the gradient of the loss
        with respect to the output, and we need to compute the gradient of the loss
        with respect to the input.
        """
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[input < 0] = 0
        return grad_input

dtype = torch.FloatTensor
# dtype = torch.cuda.FloatTensor # Uncomment this to run on GPU
# N is batch size; D_in is input dimension;
# H is hidden dimension; D_out is output dimension.
N, D_in, H, D_out = 64, 1000, 100, 10

# Create random Tensors to hold input and outputs, and wrap them in Variables.
x = Variable(torch.randn(N, D_in).type(dtype), requires_grad=False)
y = Variable(torch.randn(N, D_out).type(dtype), requires_grad=False)

# Create random Tensors for weights, and wrap them in Variables.
w1 = Variable(torch.randn(D_in, H).type(dtype), requires_grad=True)
w2 = Variable(torch.randn(H, D_out).type(dtype), requires_grad=True)

learning_rate = 1e-6
for t in range(500):
    # To apply our Function, we use Function.apply method. We alias this as 'relu'.
    relu = MyReLU.apply
    # Forward pass: compute predicted y using operations on Variables; we compute
    # ReLU using our custom autograd operation.
    y_pred = relu(x.mm(w1)).mm(w2)
    # Compute and print loss
    loss = (y_pred - y).pow(2).sum()
    print(t, loss.item())
    # Use autograd to compute the backward pass.
    loss.backward()

    # Update weights using gradient descent
    w1.data -= learning_rate * w1.grad.data
    w2.data -= learning_rate * w2.grad.data

    # Manually zero the gradients after updating weights
    w1.grad.data.zero_()
    w2.grad.data.zero_()

PyTorch: Defining New autograd Functions — PyTorch Tutorials 1.11.0+cu102 documentation

# -*- coding: utf-8 -*-
import torch
import math


class LegendrePolynomial3(torch.autograd.Function):
    """
    We can implement our own custom autograd Functions by subclassing
    torch.autograd.Function and implementing the forward and backward passes
    which operate on Tensors.
    """

    @staticmethod
    def forward(ctx, input):
        """
        In the forward pass we receive a Tensor containing the input and return
        a Tensor containing the output. ctx is a context object that can be used
        to stash information for backward computation. You can cache arbitrary
        objects for use in the backward pass using the ctx.save_for_backward method.
        """
        ctx.save_for_backward(input)
        return 0.5 * (5 * input ** 3 - 3 * input)

    @staticmethod
    def backward(ctx, grad_output):
        """
        In the backward pass we receive a Tensor containing the gradient of the loss
        with respect to the output, and we need to compute the gradient of the loss
        with respect to the input.
        """
        input, = ctx.saved_tensors
        return grad_output * 1.5 * (5 * input ** 2 - 1)


dtype = torch.float
device = torch.device("cpu")
# device = torch.device("cuda:0")  # Uncomment this to run on GPU

# Create Tensors to hold input and outputs.
# By default, requires_grad=False, which indicates that we do not need to
# compute gradients with respect to these Tensors during the backward pass.
x = torch.linspace(-math.pi, math.pi, 2000, device=device, dtype=dtype)
y = torch.sin(x)

# Create random Tensors for weights. For this example, we need
# 4 weights: y = a + b * P3(c + d * x), these weights need to be initialized
# not too far from the correct result to ensure convergence.
# Setting requires_grad=True indicates that we want to compute gradients with
# respect to these Tensors during the backward pass.
a = torch.full((), 0.0, device=device, dtype=dtype, requires_grad=True)
b = torch.full((), -1.0, device=device, dtype=dtype, requires_grad=True)
c = torch.full((), 0.0, device=device, dtype=dtype, requires_grad=True)
d = torch.full((), 0.3, device=device, dtype=dtype, requires_grad=True)

learning_rate = 5e-6
for t in range(2000):
    # To apply our Function, we use Function.apply method. We alias this as 'P3'.
    P3 = LegendrePolynomial3.apply

    # Forward pass: compute predicted y using operations; we compute
    # P3 using our custom autograd operation.
    y_pred = a + b * P3(c + d * x)

    # Compute and print loss
    loss = (y_pred - y).pow(2).sum()
    if t % 100 == 99:
        print(t, loss.item())

    # Use autograd to compute the backward pass.
    loss.backward()

    # Update weights using gradient descent
    with torch.no_grad():
        a -= learning_rate * a.grad
        b -= learning_rate * b.grad
        c -= learning_rate * c.grad
        d -= learning_rate * d.grad

        # Manually zero the gradients after updating weights
        a.grad = None
        b.grad = None
        c.grad = None
        d.grad = None

print(f'Result: y = {a.item()} + {b.item()} * P3({c.item()} + {d.item()} x)')