【他山之石】从NumPy开始实现一个支持Auto-grad的CNN框架

“他山之石，可以攻玉”，站在巨人的肩膀才能看得更高，走得更远。在科研的道路上，更需借助东风才能更快前行。为此，我们特别搜集整理了一些实用的代码链接，数据集，软件，编程技巧等，开辟“他山之石”专栏，助你乘风破浪，一路奋勇向前，敬请关注。

作者：知乎—maxweIt

地址：https://www.zhihu.com/people/maxwell-21-7

这本来是一门课的课程大作业，原本的要求是设计一个CNN推理框架，并不要求支持训练；但布置作业那天恰好我没去上课，所以把训练一起实现了。

要特别说明的是，Pico并不能被称为一个实用的框架，因为我实现的只是非常基础的功能和接口，性能上也远不如当前成熟的框架，但可以用于进一步学习了解深度学习的前向、反向计算机制，包括二维卷积的计算、部分优化器的实现等。

在实现上，基于Python和NumPy，整体框架的设计目标是在实现CNN训练、推理的同时，在上层接口和使用上接近于PyTorch，所以如果有PyTorch的基本使用基础可以更好理解。

https://github.com/maxwellzh/Pico

框架设计

Tensor

和大多数框架一样，我们将基本运算单元称为张量（Tensor）；Tensor类中包含4个基本属性：

• data: 数据真正的存储位置，使用的是NumPy的ndarray类；
• grad: 用于保存梯度，类型和data一致；
• requires_grad: 标志符，Tensor是否需要梯度；
• retain: 标志符，在计算图释放时是否释放自身

Function

算子基类，所有对Tensor的操作（例如 + − × ÷ 和卷积、池化等）都必须是Function类的派生类。如果写过PyTorch的自定义梯度计算函数，对这个应该不会陌生。Function的派生类需要实现以下方法：

• forward: 前向计算函数，静态函数（原因后面会提到）；
• backward: 反向计算函数，静态函数.

其中反向计算函数的输出必须与前向计算函数的输入参数数量相同（若参数不需要计算梯度则返回None）。特别需要注意的是，这两个函数除了计算相关参数之外，都有一个ctx输入参数，这个可以认为是一个容器，用于保存一些需要的中间变量，在方向传播时可以使用。

设想我们要实现÷（除法）运算的Function，首先理论上我们知道，前向计算为

反向计算时，对两个输入参数A和B分别求梯度，根据链式法则，有

以下是其具体实现，其中grad_out表示的就是上述公式中的  ：

_Div_类是除法的Function，可以看到我创建了一个 _Div_类的实例base_operator_div；此后，每当需要计算Tensor的除法时，只需要调用 base_operator_div(A, B)即可。

在这个例子中，计算梯度时需要用到A和B的值，因此可以看到前向计算时它们被保存在ctx中，并在反向传播时被读取出来。

回到之前提到的为什么将前向和反向计算函数实现为静态方法的问题：

• 如果实现为非静态方法，即不需要ctx参数，而是将中间变量保存在self.xxx中，那么每需要计算一次除法，就要创建一个新的_Div_类实例；
• 而实现为静态方法后，仅需要全局创建一个除法Function实例，每次需要时调用即可。

FunctionAuto-grad的实现

成熟的框架例如PyTorch、Tensorflow和mxnet等都有自动求导机制，即无需用户实现反向传播计算，这也是真正核心的部分。

我们定义一个Tracer类，并创建了一个全局唯一的tracer实例，用于记录计算图、控制反向传播等。其中包括三个成员变量：

• tensors: Tensor池，所有创建的Tensor都会被记录到其中，用C++的说法就是保存了指向Tensor的指针；
• funcs: Function池，所有调用的Function都会被记录到其中，注意这和Tensor池不一样，并不记录创建而是记录调用；
• blind: 标志符，用于控制tracer是否记录梯度.

通过Tensor池和Function池，我们可以方便地构建计算图。以下是一个简单的例子：

# 导入模块
>>> import numpy as np
>>> import pico.functional as F
>>> from pico.base import Tensor
# 创建 Tensor A 和 Tensor B
>>> A = Tensor(np.array([2.5]), requires_grad=True)
>>> B = Tensor(np.array([0.5]), requires_grad=True)
# 计算 A / B
>>> O = A/B
# 输出 O
>>> print(O)
Tensor([5.], requires_grad=True)
# O分别对A, B计算梯度
>>> O.backward()
# 输出梯度
>>> print(A.grad)
[2.]
>>> print(B.grad)
[-10.]

需要注意的是，像这样构建池记录计算图一个很重要的问题是如何清理计算图，这一块的逻辑较为复杂，具体可以参考代码。

后续有时间再补充如何基于NumPy简单实现二维卷积（Conv2d）

本文目的在于学术交流，并不代表本公众号赞同其观点或对其内容真实性负责，版权归原作者所有，如有侵权请告知删除。

“他山之石”历史文章

更多他山之石专栏文章，

请点击文章底部“阅读原文”查看

分享、点赞、在看，给个三连击呗！