动手设计实现一个深度学习框架

本文搬运于个人博客，欢迎点击 这里 查看原博文。

当前深度学习框架越来越成熟，对于使用者而言封装程度越来越高，好处就是现在可以非常快速地将这些框架作为工具使用，用非常少的代码就可以进行实验，坏处就是可能背后地实现都被隐藏起来了。在这篇文章里笔者将带大家一起用 Python 从头设计和实现一个轻量级的、易于扩展的深度学习框架 tinynn，希望对大家了解深度学习原理与深度学习框架有一定的帮助。

本文首先会从深度学习的流程开始分析，对神经网络中的关键组件抽象，确定基本框架；然后再对框架里各个组件进行代码实现；最后基于这个框架实现了一个 MNIST 分类的示例。

目录

• 组件抽象
• 组件实现
• 整体结构
• MNIST 例子
• 总结
• 附录

组件抽象

首先考虑神经网络运算的流程，神经网络运算主要包含训练 training 和预测 predict （或 inference） 两个阶段，训练的基本流程是：输入数据 -> 网络层前向传播 -> 计算损失 -> 网络层反向传播梯度 -> 更新参数，预测的基本流程是 输入数据 -> 网络层前向传播 -> 输出结果。从运算的角度看，主要可以分为三种类型的计算：

1. 数据在网络层直接的流动
   前向传播和反向传播可以看做是张量 Tensor（多维数组）在网络层之间的流动（前向传播流动的是输入输出，反向传播流动的是梯度），每个网络层会进行一定的运算，然后将结果输入给下一层

2. 计算损失
   衔接前向和反向传播的中间过程，定义了模型的输出与真实值之间的差异，用来后续提供反向传播所需的信息

3. 参数更新
   使用计算得到的梯度对网络参数进行更新的一类计算

基于这个三种类型，我们可以对网络的基本组件做一个抽象

• tensor 张量，这个是神经网络中数据的基本单位
• layer 网络层，负责接收上一层的输入，进行该层的运算，将结果输出给下一层，由于 tensor 的流动有前向和反向两个方向，因此对于每种类型网络层我们都需要同时实现 forward 和 backward 两种运算
• loss 损失，在给定模型预测值与真实值之后，该组件输出损失值以及关于最后一层的梯度（用于梯度回传）
• optimizer 优化器，负责使用梯度更新模型的参数

然后我们还需要一些组件把上面这个 4 种基本组件整合到一起，形成一个 pipeline

• net 组件负责管理 tensor 在 layer 之间的前向和反向传播，同时能提供获取参数、设置参数、获取梯度的接口
• model 组件负责整合所有组件，形成整个 pipeline。即 net 组件进行前向传播 -> loss 组件计算损失和梯度 -> net 组件将梯度反向传播 -> optimizer 组件将梯度更新到参数。

基本的框架图如下图

组件实现

按照上面的抽象，我们可以写出整个流程代码如下。首先定义 net，net 的输入是多个网络层，然后将 net、loss、optimizer 一起传给 model。model 实现了 forward、backward 和 apply_grad 三个接口分别对应前向传播、反向传播和参数更新三个功能。

# define model
net = Net([layer1, layer2, ...])
model = Model(net, loss_fn, optimizer)

# training
pred = model.forward(train_X)
loss, grads = model.backward(pred, train_Y)
model.apply_grad(grads)

# inference
test_pred = model.forward(test_X)
复制代码

接下来我们看这里边各个部分分别如何实现。

• tensor

  tensor 张量是神经网络中基本的数据单位，我们这里直接使用 numpy.ndarray 类作为 tensor 类的实现（numpy 底层使用了 C 和 Fortran，并且在算法层面进行了大量的优化，运算速度也不算慢）

• layer

  上面流程代码中 model 进行 forward 和 backward，其实底层都是网络层在进行实际运算，因此网络层需要有提供 forward 和 backward 接口进行对应的运算。同时还应该将该层的参数和梯度记录下来。先实现一个基类如下

  # layer.py
  class Layer(object):
      def __init__(self, name):
          self.name = name
          self.params, self.grads = None, None
  
      def forward(self, inputs):
          raise NotImplementedError
  
      def backward(self, grad):
          raise NotImplementedError
  复制代码

  最基础的一种网络层是全连接网络层，实现如下。forward 方法接收上层的输入 inputs，实现 $wx+b$ 的运算；backward 的方法接收来自上层的梯度，计算关于参数 $w, b$ 和输入的梯度，然后返回关于输入的梯度。这三个梯度的推导可以见附录，这里直接给出实现。w_init 和 b_init 分别是参数 weight 和 bias 的初始化器，这个我们在另外的一个实现初始化器中文件 initializer.py 去实现，这部分不是核心部件，所以在这里不展开介绍。

  # layer.py
  class Dense(Layer):
      def __init__(self, num_in, num_out,
                   w_init=XavierUniformInit(),
                   b_init=ZerosInit()):
          super().__init__("Linear")
  
          self.params = {
              "w": w_init([num_in, num_out]),
              "b": b_init([1, num_out])}
  
          self.inputs = None
  
      def forward(self, inputs):
          self.inputs = inputs
          return inputs @ self.params["w"] + self.params["b"]
  
      def backward(self, grad):
          self.grads["w"] = self.inputs.T @ grad
          self.grads["b"] = np.sum(grad, axis=0)
          return grad @ self.params["w"].T
  复制代码

  同时神经网络中的另一个重要的部分是激活函数。激活函数可以看做是一种网络层，同样需要实现 forward 和 backward 方法。我们通过继承 Layer 类实现激活函数类，这里实现了最常用的 ReLU 激活函数。func 和 derivation_func 方法分别实现对应激活函数的正向计算和梯度计算。

  # layer.py
  class Activation(Layer):
  		"""Base activation layer"""
      def __init__(self, name):
          super().__init__(name)
          self.inputs = None
  
      def forward(self, inputs):
          self.inputs = inputs
          return self.func(inputs)
  
      def backward(self, grad):
          return self.derivative_func(self.inputs) * grad
  
      def func(self, x):
          raise NotImplementedError
  
      def derivative_func(self, x):
          raise NotImplementedError
  
  class ReLU(Activation):
  		"""ReLU activation function"""
      def __init__(self):
          super().__init__("ReLU")
  
      def func(self, x):
          return np.maximum(x, 0.0)
  
      def derivative_func(self, x):
          return x > 0.0
  复制代码

• net

  上文提到 net 类负责管理 tensor 在 layer 之间的前向和反向传播。forward 方法很简单，按顺序遍历所有层，每层计算的输出作为下一层的输入；backward 则逆序遍历所有层，将每层的梯度作为下一层的输入。这里我们还将每个网络层参数的梯度保存下来返回，后面参数更新需要用到。另外 net 类还实现了获取参数、设置参数、获取梯度的接口，也是后面参数更新时需要用到

  # net.py
  class Net(object):
      def __init__(self, layers):
          self.layers = layers
  
      def forward(self, inputs):
          for layer in self.layers:
              inputs = layer.forward(inputs)
          return inputs
  
      def backward(self, grad):
          all_grads = []
          for layer in reversed(self.layers):
              grad = layer.backward(grad)
              all_grads.append(layer.grads)
          return all_grads[::-1]
  
      def get_params_and_grads(self):
          for layer in self.layers:
              yield layer.params, layer.grads
  
      def get_parameters(self):
          return [layer.params for layer in self.layers]
  
      def set_parameters(self, params):
          for i, layer in enumerate(self.layers):
              for key in layer.params.keys():
                  layer.params[key] = params[i][key]
  复制代码

• loss

  上文我们提到 loss 组件需要做两件事情，给定了预测值和真实值，需要计算损失值和关于预测值的梯度。我们分别实现为 loss 和 grad 两个方法，这里我们实现多分类回归常用的 SoftmaxCrossEntropyLoss 损失。这个的损失 loss 和梯度 grad 的计算公式推导进文末附录，这里直接给出结果：

  多分类 softmax 交叉熵的损失为

  $$ J_{CE}(y, \hat{y}) = -\sum_{i=1}^N \log \hat{y_i^{c}} $$

  © 版权声明

  文章版权归作者所有，未经允许请勿转载。
  THE END
  人工智能