pytorch的自定义拓展之（一）——torch.nn.Module和torch.autograd.Function

前言：pytorch的灵活性体现在它可以任意拓展我们所需要的内容，前面讲过的自定义模型、自定义层、自定义激活函数、自定义损失函数都属于pytorch的拓展，这里有三个重要的概念需要事先明确。要实现自定义拓展，有两种方式，

（1）方式一：通过继承torch.nn.Module类来实现拓展。这也是我们前面的例子中所用到的，它最大的特点是以下几点：

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

（2）方式二：通过继承torch.nn.Function类来实现拓展。它最大的特点是：

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

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

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

pytorch中有着自动求导机制，当然这针对的仅仅是torch里面所定义的一些函数，我们知道torch.nn.functional是专门为神经网络所定义的函数集合），如果我们有时候需要进行的操作是nn.functional中没有提供，甚至是torch里面也没有提供的，那怎么办呢？当然我们可以使用一些基本的pytorch函数来进行组装，另外我们也可以使用numpy或scipy三方库中的方法实现。这个时候

由于pytorch不再提供自动求导机制，就要自己定义实现前向传播和反向传播的计算过程了。

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

1.1 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):

当然这里没有列举完全，他还有一些属性和方法是定义在Function的父类里面的，这里就不再一一列举了。

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

（1）torch.autograd.Function类实际上是某一个操作函数的父类，一个操作函数必须具备两个基本的过程，即前向的运算过程和反向的求导过程，

（2）torch.nn.Module类实际上是对torch.xxxx以及torch.nn.functional.xxxx这些函数的包装组合，而torch.xxxx和torch.nn.functional.xxxx都是实现了autograd.Function类的两个基本功能（前向运算和反向传播），如果是我们需要的某一个功能torch.xxxx和torch.nn.functional里面都没有，也不能通过组合得到，这就需要定义新的操作函数，这个函数就需要继承自autograd.Function类，重写前向运算和反向传播。（注意体会这段话）

（3）很显然，nn.Module更加高层，而autograd.Function更加底层，其实从名字中也能看出二者的区别，Module是针对模块的，即神经网络中的层、激活层、损失函数、网络模型等等，而Function是针对函数的，针对的是一些需要自己定义的函数而言的。如果某一个函数my_function继承自Function类，实现了这个类的forward和backward方法，那么我依然可以用nn.Module对这个自定义的的函数my_function进行包装组合，因为此时my_function跟torch.xxxx和torch.nn.functional.xxxx里面的函数已经具备了等同的地位了。（注意体会这段话），可以这么说，Module不仅包括了Function，还包括了对应的参数，以及其他函数与变量，这是Function所不具备的。

（4）那为什么Function类也可以定义一个神经网络呢？

在官网的例子中，我们常常看见下面这样的定义：

class MyReLU(torch.autograd.Function):
    def forward(self, input_):  

    def backward(self, grad_output):

input_ = Variable(torch.linspace(-3, 3, steps=5)) # 定义输入
my_relu=MyReLU()   # 构建模型
output_ = my_relu(input_)

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

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

1.2 总结：

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

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

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

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

鉴于这个类确实是比较底层，正在使用的时候经常遇见我找不到的原因，所以本文只列举较为简单的情况，即不使用torch之外的三方库（numpy、scipy等，由于numpy和scipy函数是不支持backward的，所以在使用的时候涉及到ndarray与tensor之间的转换，常常出错），另外也暂时不涉及向量对向量的求导，仅仅涉及标量对标量和标量对向量求导，这里可以参考我的前面一篇文章：pytorch自动求导Autograd系列教程（一）

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
import numpy as np

# 定义一个继承了Function类的子类，实现y=f(x)的正向运算以及反向求导
class sqrt_and_inverse(torch.autograd.Function):
    '''
    forward和backward可以定义成静态方法，向定义中那样，也可以定义成实例方法
    '''
    # 前向运算
    def forward(self, input_x,input_y): 
       '''
       self.save_for_backward(input_x,input_y) ,这个函数是定义在Function的父类_ContextMethodMixin中 
            它是将函数的输入参数保存起来以便后面在求导时候再使用，起前向反向传播中协调作用    
       ''' 
        self.save_for_backward(input_x,input_y)                  
        output=torch.sqrt(input_x)+torch.reciprocal(input_x)+2*torch.pow(input_y,2)
        return output                              
                                         
    def backward(self, grad_output):                             
        input_x,input_y=self.saved_tensors  # 获取前面保存的参数,也可以使用self.saved_variables
        grad_x = grad_output *(torch.reciprocal(2*torch.sqrt(input_x))-torch.reciprocal(torch.pow(input_x,2)))
        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 sqrt_and_inverse()(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
import numpy as np

class sqrt_and_inverse(torch.autograd.Function):
                                 
    def forward(self, input_x):  #input_x是一个tensor，不再是一个标量
        self.save_for_backward(input_x)                  
        output=torch.sum(torch.sqrt(torch.pow(input_x,2)-1)) # 函数z
        return output                             
                                         
    def backward(self, grad_output):                                 
        input_x,=self.saved_tensors  # 获取前面保存的参数,也可以使用self.saved_variables  #input_x前面的逗号是不能丢的
        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()(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进行拓展的一般模板

class My_Function(Function):
 def forward(self, inputs, parameters):
        self.saved_for_backward = [inputs, parameters]
        # output = [对输入和参数进行的操作，其实就是前向运算的函数表达式]
        return output

 def backward(self, grad_output):
        inputs, parameters = self.saved_tensors # 或者是self.saved_variables
        # grad_inputs = [求函数forward(input)关于 parameters 的导数，其实就是反向运算的导数表达式] * grad_output
        return grad_input

自定义类的包装

# 包装自定义的My_Function有几种方法，通过方法包装，通过一个类包装都可以
# 这里就展示使用一个方法包装
# 这样使得看起来更加自然，因为Function的作用就是实现一个自定义方法的
def my_function(inputs):
    return My_Function()(inputs) # 一定要是对象调用
    
'''注意事项：
需要注意的是，这里一定要使用对象调用，否则虽然也能够求出倒数结果，但实际上跟我自己定义backward函数就没啥关系了
如果使用 return My_Function().forward(inputs)
这是不行的，虽然结果正确，后面会分析
'''

然后我们就可以将我们自己所定义的方法（也就是继承自Function的类）像pytorch自己定义的方法那样去使用了。

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

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

包装函数里面一定要使用return My_Function()(inputs) 即对象调用，而不能使用，return My_Function().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前面的逗号是不能丢的
    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()(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().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(self, 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类型，当是标量求导的时候，它是一个标量值，当是向量求导的时候，它是一个和求导向量同维度的向量。具体可参见前面的文章：pytorch自动求导Autograd系列教程（一）

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

print('开始前向传播')
z=sqrt_and_inverse_func(x,y)  
print(z) 
print(z.grad_fn)  
'''运行结果为：
开始前向传播
tensor(10.0654, grad_fn=)
<__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)  # 返回  表示是由Add加法创造得到的Y

 

三、autograd.Function的更多应用

参见下一篇文章