使用numpy搭建自己的深度学习框架（二）

注：本系列搭建的深度学习框架名称叫numpyflow，缩写nf，用以熟悉目前主流的深度学习框架的基础和原理。本系列的目标是使用nf可以训练resnet。
开源地址：RanFeng/NumpyFlow

简介

上一节自动求导（基础知识）中我们介绍了自动求导的必备要素，其中之一就是要将一些基本的操作，比如加、减、乘、除、乘方等封装成一个个基本的算子，并首先写好对应的梯度计算方法，此处给一段代码作为例子：

class Add:

    def __call__(self, a, b):
        self.variables = (a, b)
        out = a.data + b.data
        return out
        
    def backward(self, grad, **kwargs):
        self.variables[0].backward(grad)
        self.variables[1].backward(grad)

这是加法算子，我们重写了内置的__call__方法使得计算的调用更加简单快捷，利用backward方法完成反向传播的梯度计算和传播。

对于其他的算子我们都可以用这个方法进行封装。同时我们需要再创建一个数据载体，类似于torch中的Tensor，我们也需要一个Tensor类。

add_op = None
class Tensor:
    def __init__(self, data=None):
        self.data = data.copy()
        self.grad = np.zeros_like(self.data, dtype=np.float64)
	def __add__(self, other):
		global add_op
		add_op = Add()
        return add_op(self, other)

上面的Tensor是个简易的Tensor类，包含了data域，接受numpy.ndarray作为输入，所以Tensor其实就是一个封装了的numpy.ndarray类，同时，我们需要grad域，用来保存上一节所计算的梯度，grad初始化为0，并且保证shape与data一致。同时我们还重写了内置的__add__方法，这样就是重载了+这个符号了。注意，__add__方法中，我们保留了add_op，这是为了保证反向传播时，能够追踪变量用的，这个在接下来的例子中会解释到。

经过上面一层的简单包装，我们的Tensor类就完成了加法操作的前向传播，这个很简单，接下来，我们要完成加法操作的反向传播，也就是自动求导了。

我们给Tensor类加一个backward方法:

class Tensor:
	...
    def backward(self, grad=None):
        self.grad += grad

好了，这样加法的反向传播也可以搞定了。我们举个例子来分析一下，比如：
$$\begin{gathered} x=3,y=4 \\ z=x+y \end{gathered}$$
求解$\frac{\partial z}{\partial x}(3,4)$和$\frac{\partial z}{\partial y}(3,4)$的值分别是多少，这个简单的问题，可以直接看出答案分别是1和1。我们现在来分析一下，在Tensor的自动求导中发生了什么。

import numpy as np
add_op = None
x = np.array([3],dtype=np.float64)
y = np.array([4],dtype=np.float64)
x = Tensor(x)	# 3
y = Tensor(y)	# 4
z = x + y		# 7

上面一段代码完成了前向的计算，同时也建立了计算图。计算图建立如下：
其中的加法op就是代码中的add_op，在这段程序中是全局的变量，并且此处的z的类型是numpy.ndarray，并不是Tensor，所以，这只是个demo，只是用来理解自动计算梯度的过程的，具体的代码请参考前面提到的github链接。

好了，计算图在前向计算的时候已经构建好了，那接下来就是反向计算梯度了，我们使用如下语句调用反向计算梯度。

add_op.backward(np.ones_like(z, dtype=np.float64))
print(x.grad, y.grad)	# 1 1

我们利用add_op完成了$z=x+y$的梯度计算，完整代码如下：

import numpy as np

class Add:
    def __call__(self, a, b):
        self.variables = (a, b)
        out = a.data + b.data
        return out

    def backward(self, grad, **kwargs):
        self.variables[0].backward(grad)
        self.variables[1].backward(grad)

class Tensor:
    def __init__(self, data=None):
        self.data = data.copy()
        self.grad = np.zeros_like(self.data, dtype=np.float64)

    def __add__(self, other):
        global add_op
        add_op = Add()
        return add_op(self, other)

    def backward(self, grad=None):
        self.grad += grad


if __name__ == '__main__':
    add_op = None
    x = np.array([3], dtype=np.float64)
    y = np.array([4], dtype=np.float64)
    x = Tensor(x)  # 3
    y = Tensor(y)  # 4
    z = x + y	   # 7
    print(z, add_op)	# [7.] <__main__.Add object at 0x112550a58>
    add_op.backward(np.ones_like(z, dtype=np.float64))
    print(x.grad, y.grad)	# [1.] [1.]

复杂的功能由简单的组成，我们下一章继续完善这个Tensor类和算子类，使得最终能完成这样一个复杂函数的运算和反向传播：

x = np.random.random([2,4,6,3,4])
y = np.random.random([2,4,6,3,4])
z = np.random.random([2,4,1,1,4])
def func(x,y,z):
    f0 = (x[1,0].T * y[0,1].T).T * z * x
    f1 = f0 * (x + y + z) * y * y * y * (y+z)
    f2 = y[0,3] + x[0,2]
    f3 = y * y - z
    f4 = z - x
    f5 = -x.flatten() + y.flatten() - (x*z).flatten() * 2.0
    f6 = f1[1,3] + f1[0,3] * f2 - z[0,1] ** 2.2
    f7 = f3 + f4 + f6
    f8 = f7 - f3 + f4 * 3.6
    f9 = f8.flatten() / f5 + f7.flatten()
    f10 = -f9 * f5
    f11 = ((x*z) @ x.transpose(3, 4) @ y.permute(0,4,2,3,1)).transpose(0,4)
    f12 = f11.transpose(3,4).flatten() * 5.0 ** x.transpose(1,4).flatten() / y.flatten() * (x/z).flatten() + 2.0
    f13 = f10.reshape(f11.shape) * f11 / f12.reshape(f11.shape)
    f14 = (x.transpose(3,4) @ y).permute(0,2,4,3,1) @ f13.permute(4,2,0,1,3)
    f15 = f14.sum() * f14.mean((0,2))
    return f15