前言：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=<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 0x000001F2E56D28D0> 表示是由Add加法创造得到的Y
三、autograd.Function的更多应用
参见下一篇文章
