DLsys - Needle 源码剖析

NEEDLE: Necessary Elements of Deep LEarning

Needle是一个微型的深度学习框架，是一个由 CMU 10-414/714 课程而提出的实验性质的框架。

该深度学习框架具备一个深度学习框架所需的基本要素，比如：以 CPU 和 CUDA 为后端的张量运算，动态计算图，模块化的神经网络编程组件，自动微分工具，经典模型实现等等。

笔者发现互联网上关于 Needle 的相关资料非常少，想着自己写一下源码剖析。

Needle 概览

Needle 本质上是一个实验课的大作业，这里我给出一个完成品的参考实现

Needle 的代码框架如下图所示

Needle
├── 3rdparty        # 一些外部依赖库，如 pybind11
├── CMakeLists.txt  # Cmake 编译脚本文件
├── LICENSE
├── README.md
├── apps            # Needle 的案例演示
├── data            # 数据存放处
├── figures         # 相关图片
├── python          # Needle 的 Python 实现
└── src             # CPU 和 CUDA 加速实现

接下来介绍一下如何使用 needle：

所需依赖：cmake、cuda、pybind11

• CMake: 一个跨平台的构建系统，帮助配置和生成项目所需的构建文件，用于编译 Needle 项目。

• CUDA: NVIDIA 提供的并行计算框架，支持在 GPU 上进行高效的计算，加速 Needle 项目中涉及的 GPU 任务。

• Pybind11: 一个用于将 C++ 代码绑定到 Python 的库，允许在 Needle 项目中无缝集成和调用 C++ 与 Python 代码。

环境配置

首先需要准备 cmake，在 Linux 操作系统下你可以通过 apt 包管理器来快速安装 cmake：

sudo apt install cmake

CUDA 的安装可以上 Nvidia 的官方指引查看，相关链接

你可以通过运行 nvcc --version 指令来查看 CUDA 环境是否配置成功。

项目需要下载 pybind11 至 3rdparty/pybind11 路径下，参考命令为：

cd 3rdparty
git clone https://github.com/pybind/pybind11.git -b stable

编译 CPU 和 CUDA 的加速算子

cd needle
mkdir build && cd build
cmake .. && make -j `nproc`

此部分的指令用于编译 CPU 和 CUDA 的底层算子实现，Needle 会通过 pybind11 在 Python 调用编译后的运行库。

此部分更细致的分析将在下文的 CUDA 加速部分涉及。

打包为 Python 第三方库

cd python
python setup.py install --user

此部分的指令用于将 Needle 封装为 Python 第三方库，并通过 Pip 安装。

此部分更细致的分析将在下文的 项目构建脚本部分涉及。

接下来笔者将从以下几个部分来拆解并分析 Needle:

• 自动微分：Needle 的自动微分部分通过将操作封装为 Op 类，利用 Value 类和拓扑排序实现前向传播和反向传播，自动计算和存储梯度，从而支持高效的梯度计算和优化。

• Tensor：#TODO!

• CUDA 加速：#TODO!

• 神经网络组件：#TODO!

• 数据加载组件：#TODO!

• 参数优化算法：#TODO!

• 项目构建脚本：#TODO!

• 案例演示：#TODO!

自动微分

这部分的代码实现主要是在 python/needle/autograd.py 和 python/needle/ops.py 文件中，主要实现了专门用于自动微分的数据结构和运算重载。

概览

我们都知道 Tensor 是深度学习框架常用的数据结构，Tensor 能够实现自动微分。为了实现自动微分，Needle 并没有直接实现 Tensor，而是先把自动微分相关的操作独立抽象为一个 Value 类，当 Value 实现自动微分后，一切继承自 Value 的类(比如 Tensor)也便直接拥有了自动微分的功能。

Needle 还对 Value 类的一切基本运算(比如加减乘除)进行重载，使得所有运算操作都是 Op 的子类。Op 类也是 Needle 定义的自动微分相关的基础类型，这样子做有两个用处：

1. 统一接口：通过对基本运算的重载，所有的操作都遵循统一的接口。每个继承自 Op 的子类代表一种特定的运算(如加法、乘法等)，并且都实现了 Op.forward() 和 Op.backward() 两个方法。这样在实现新的运算时，只需继承 Op 类并实现相应的方法，从而保证了代码的一致性和可扩展性。

2. 链式法则自动化：每当进行运算时，Needle 会记录每次 Op 运算操作的信息，包括但不限于参与运算的元素、本次运算的类型以及运算结果元素，这将构建一个计算图。随后，通过调用 Op.backward() 方法，Needle 能够依据刚刚建立的计算图自动计算出所有相关变量的梯度，从而实现自动微分。

接下来说一下梯度是如何计算的，当我们实现一个操作比如加法时，我们需要实现这个操作对应的 forward() 和 backward() 函数，前者用于前向传播计算，后者用于反向传播计算。在 backward() 我们需要为所有操作写好梯度计算的规则，当需要进行反向传播时，系统会沿着这个计算图，从输出开始逆向遍历每个操作节点，并按照定义好的 backward() 规则，逐步计算出每个输入变量的梯度。

自动微分的 “自动” 部分正体现在这里：我们需要为每种基本操作定义梯度计算规则，前向计算时，Needle 会根据操作顺序自动构建计算图(表现为记下每次操作涉及的元素和类型)，并在反向传播时自动应用这些梯度规则，而不需要手动计算每一步的导数。基础的操作会组合成复杂的操作，这使得复杂的梯度计算变得更加高效和可靠，尤其是在涉及多个变量和复杂操作的深度学习模型中。

Value 类

Value 的成员变量

class Value:
  op: Optional[Op]
  inputs: List["Value"]
  cached_data: Array
  requires_grad: bool

Value 类一共有四个成员变量，解读如下：

• op：类型为 Optional[Op]，用于表示此 Value 变量的运算来源，如果是加法表示此变量是一次加法运算的结果，如果为 None 则表示此变量为直接创建得到。
• inputs：类型为 List["Value"]，用于记录得到此 Value 变量的运算父母元素，如果为空列表则表示此变量为直接创建得到。
• cached_data：类型为 Array，即此 Value 变量存储的数据，通常为矩阵。
• requires_grad：类型为 bool，即表示此 Value 变量在计算时是否参与计算图的构建，默认为 True，若此值为 False，则此变量的任何运算结果都不会记录此变量的信息。

可以发现每个 Value 变量都会记录其运算父母元素和运算类型，下文我们称没有运算父母元素的变量为叶子变量，叶子变量是反向梯度计算的终点。

接下来我们看一下 Value 类的一些用于自动微分的函数。

Value.init()

class Value:
    def _init(
        self,
        op: Optional[Op],
        inputs: List["Tensor"],
        *,
        num_outputs: int = 1,
        cached_data: List[object] = None,
        requires_grad: Optional[bool] = None
    ):
        global TENSOR_COUNTER
        TENSOR_COUNTER += 1
                
        if requires_grad is None:
            requires_grad = any(x.requires_grad for x in inputs)
        self.op = op
        self.inputs = inputs
        self.num_outputs = num_outputs
        self.cached_data = cached_data
        self.requires_grad = requires_grad
        return

这段代码主要用于属性赋值，用来构造 Value 类，实现逻辑非常简单就不赘述了。

Value.is_leaf()

class Value:
    def is_leaf(self):
        return self.op is None

这个函数用来判断此变量是否为叶子变量，逻辑非常简单，判断此变量是否记录了运算父母信息即可。

Op 类及其重载

Op 类的函数

class Op:
    def __call__(self, *args):
        raise NotImplementedError()

    def compute(self, *args: Tuple[Array]):
        raise NotImplementedError()

    def gradient(
        self, out_grad: "Value", node: "Value"
    ) -> Union["Value", Tuple["Value"]]:
        raise NotImplementedError()

Op 类是一个抽象基类，定义了在自动微分框架中每个操作必须实现的核心接口。这个类的设计使得所有具体的操作(如加法、乘法等)都必须继承 Op 类并实现以下三个关键方法：

1. **Op.__call__()**：这是一个抽象方法，用于定义操作的调用行为。

2. **Op.compute()**：这个方法用于执行具体的前向传播计算。

3. **Op.gradient()**：这个方法用于计算反向传播时的梯度。

通过将这三个方法定义为抽象方法，Op 类确保了所有的操作都有统一的接口，并且每个具体操作都必须显式实现这些方法，这样可以确保自动微分系统在前向传播和反向传播中正确地执行每个操作的逻辑。

Op 重载(以乘法为例)

Needle 重载了非常多的操作，对于大部分的操作而言逻辑是相通的，这里我们以乘法的运算重载为例，解读一下基类 Op 的重载和使用方法。

class EWiseMul(Op):
    def compute(self, a: Array, b: Array):
        return a * b

    def gradient(self, out_grad: Tensor, node: Tensor):
        lhs, rhs = node.inputs
        return out_grad * rhs, out_grad * lhs

Needle 重载了许多操作，对于大部分操作来说，逻辑是一致的。以 EWiseMul(逐元素乘法)的运算重载为例，我们可以解读基类 Op 的重载和使用方法。

EWiseMul 类继承自 Op，并重载了两个关键方法：

1. compute 方法：

   • compute(self, a: Array, b: Array): 这个方法执行前向传播的实际计算。对于逐元素乘法，它直接返回输入数组 a 和 b 的逐元素相乘结果。

2. gradient 方法：

   • gradient(self, out_grad: Tensor, node: Tensor): 该方法用于反向传播，计算梯度。对于逐元素乘法，输出的梯度 out_grad 分别乘以输入中的另一个张量（lhs 或 rhs），返回相应的梯度。这利用了链式法则，实现了自动微分。

这一示例展示了 Op 类的具体操作是如何通过重载 compute 和 gradient 方法来实现前向传播和反向传播的。

反向传播函数

前文已经说明了 Value 和 Op 类，接下来将目光放在反向传播的函数上。

topo_sort_dfs()

def topo_sort_dfs(node, visited, topo_order):
    """后序 DFS 遍历，用于拓扑排序。"""
    if hash(node) not in visited:
        visited.append(hash(node))
        for node_input in node.inputs:
            topo_sort_dfs(node_input, visited, topo_order)
        topo_order.append(node)

递归遍历节点的输入，确保所有后继节点在当前节点之前被处理，生成拓扑排序。

find_topo_sort()

def find_topo_sort(node_list: List[Value]) -> List[Value]:
    """返回节点的拓扑排序列表。"""
    visited = []
    topo_order = []
    for node in node_list:
        if hash(node) not in visited:
            topo_sort_dfs(node, visited, topo_order)
    return topo_order

通过对给定节点列表的 DFS 遍历来实现。确保每个节点在其所有前驱节点之后被添加到排序列表中，从而实现拓扑排序。

compute_gradient_of_variables()

def compute_gradient_of_variables(output_tensor, out_grad):
    """计算每个节点的梯度，并存储在每个变量的 grad 字段中。"""
    # 从输出节点到每个节点的梯度贡献的映射
    node_to_output_grads_list: Dict[Tensor, List[Tensor]] = {}
    # 初始化输出节点的梯度
    node_to_output_grads_list[output_tensor] = [out_grad.detach()]
    # 根据输出节点的逆拓扑排序遍历计算图
    reverse_topo_order = list(reversed(find_topo_sort([output_tensor])))

    for idx, node in enumerate(reverse_topo_order):
        # 计算当前节点的梯度
        node.grad = sum_node_list(node_to_output_grads_list[node])

        if node.op:  # 不是原始输入节点
            # 计算输入节点的梯度
            grads = node.op.gradient_as_tuple(node.grad, node)
            for grad, input_node in zip(grads, node.inputs):
                if input_node not in node_to_output_grads_list:
                    node_to_output_grads_list[input_node] = [grad.detach()]
                else:
                    node_to_output_grads_list[input_node].append(grad.detach())

首先维护每个节点及其对应的梯度列表 node_to_output_grads_list。初始化时，将 output_tensor 的梯度设置为传入的 out_grad。通过 find_topo_sort() 获取拓扑排序，使用拓扑排序的逆序遍历计算图，以确保在计算梯度时每个节点的输入节点已经被处理。对于每个节点，首先计算其梯度，然后更新其输入节点的梯度。

这些函数共同作用实现了反向传播中的梯度计算过程。compute_gradient_of_variables() 函数负责从输出节点开始，按照拓扑排序的逆序计算每个节点的梯度，而 find_topo_sort() 和 topo_sort_dfs() 函数则用于生成计算图的拓扑排序。

反向传播通过计算梯度并将其存储在每个节点的 grad 字段中，实现了自动微分。通过拓扑排序，确保了在计算每个节点的梯度时，其所有前驱节点的梯度已经被计算完毕。

Tensor

#TODO!

CUDA 加速

#TODO!

神经网络组件

#TODO!

数据加载组件

#TODO!

参数优化算法

#TODO!

项目构建脚本

#TODO!

案例演示

#TODO!