CMU DLSys Course Homework 1: Implementation and Reflection

Jan. 08, 2025 - Jan. XX, 2025 · Qiyao Wang · Topic #MLSys

本周看完了课程的 Lecture 3/4/5，这几节课程是对自动求导（Automatic Differentiation）的讲解，其中反向传播、计算图是重点。本文针对 Homework 1 进行实践与思考。

Introduction of Homework 1

通过 Homework 1 自己实践构建一个 needle（Necessary Elements of Deep Learning） 库，其中本文最关键的是自动求导框架的构建，并在 Homework 0 中的双层神经网络上进行应用。

本文对于 needle 的构建是建立在 numpy 库和 CPU 基础上的，后续的作业也会实现相应的底层多维数组以及应用到 GPU 节点上。

Homework 1 的 needle 库包含两个重要的文件，即 autograd.py 和 ops_mathematic.py，其中前者定义了计算图框架以及自动求导框架，后者包含在计算图中会用到的运算符（compute function），及运算符对应的求导函数（gradient function）。

Basic Concept of Needle

本文中应用 needle 库中最重要的几个类包括张量 Tensor 和其继承的父类 Value，以及对应的张量运算符 TensorOp 和其继承的父类 Op。下面对这几个类进行详细解释

  • Value (the most abstract class)

    Value 类是抽象程度最高的类，其实例化的每一个对象都作为计算图中的一个节点（甚至可直接认为其对象代表一个计算图，因为可通过节点进行轨迹追溯）。其中包含四个重要的部分：op: Optional[Op]，该语句定义了节点 Value 所链接的其他节点进行的操作，如假设 $v_3=v_1+v_2$，则 op 中存储的是 ops_mathematic.py 中定义的 “+”（EWiseAdd 逐元素加法），其中 Optional 代表 op 的数据类型要么是 Op Class，要么是 None；inputs: List["Value"]，其中包含了计算图中与当前节点相关的节点，在 $v_3=v_1+v_2$ 的例子中，$v_3$ 的 inputs 则为 $[v_1,v_2]$，通过 op 和 inputs 可以有效地获得对应计算图的完整轨迹；cached_data: NDArray，则是该节点存储的具体的值，在本作业中还是使用 numpy 中的多维数组来简化作业；requires_grad: bool，则是对该节点是否需要求导，当值为 False 时，表示该节点无需求导，无需进行梯度下降的优化。

  • Tensor (the subclass of value and the interface with user)

    Tensor 类是 Value 的子类，并且对应一个实际的张量节点（计算图中的一个多维数组），计算图中实际应用 Tensor 类而非 Value 类。

  • Op (Operations)

    Op 类用于构建计算图中包含的具体的运算符，每一个运算符需要定义计算的前馈函数 compute() 和该运算符对应的求导过程 gradient()，用于后续的反向传播。

  • TensorOp (the subclass of Op)

    TensorOp 类是 Op 的子类，它继承了 Op 中的所有操作符，并且能够返回 Tensor 类型的结果。

在实际运行 needle 库时，需要使用的是 Tensor 类，而 TensorOp 等类都内置于 Tensor 中，以 $c = a + b$ 为例

import needle as ndl
a = ndl.Tensor([1], dtype="float32")
b = ndl.Tensor([2], dtype="float32")
c = a + b # c: Tensor([3], dtype="float32")

上述代码的实际运算过程如下：

  1. Tensor.__add__ 调用 Tensor 类中的 "__add__" 方法，由于运算符两侧均为 Tensor，则实际运行的是 Op 中的逐元素加法 EWiseAdd()；
  2. TensorOp.__call__；
  3. Tensor.make_from_op(op: Op, inputs: List["Value"]) 首先为 c 实例化了一个新的 Tensor 对象，而非在 a/b 上进行修改，将对应的运算符和参与运算的节点传入新创建的对象 c 中进行初始化，其中需要对 LAZY_MODE 进行判断，最终返回新的 Tensor 对象 c；
  4. 在返回之前，经过了 Value.realize_cached_data() 或 Tensor.detach() 的计算，二者区别（及 LAZY_MODE）将在后面详细说明；
  5. 进入 Op 类的前馈计算函数中 EWiseAdd.compute 才能得到第 3 步中返回的新 Tensor 对象 c。

LAZY_MODE and Detach Function

如下方的代码块所示，为 Value 类中的成员函数 make_from_op()，其中包含了这两个元素

@classmethod
def make_from_op(cls, op: Op, inputs: List["Value"]):
  value = cls.__new__(cls) # create a new object of Value Class
  value._init(op, inputs) # initialize object
  if not LAZY_MODE:
    if not value.requires_grad:
      return value.detach()
    value.realize_cached_data()
  return value

LAZY_MODE 用来控制计算立即执行还是延迟执行。当设置 ndl.autograd.LAZY_MODE = True 时，运行 $c = a + b$ 时并不会立刻进行 realize_cached_data() 的计算，此时的 c.cached_data is None == True，这样只构建计算图但不进行计算，能够有效地优化内存占用，当需要使用 $c$ 的数值时，可以利用 c.data() 来通过前馈计算图获得该节点的 cache_data 数值。

Detach function 用来控制节点是否与计算图链接。如下方的代码所示，虽然看起来简单应该不会占用太多内存，但是由于 Tensor 的计算实质上是对计算图的构建，如图 1 所示，在新节点建立时，背后其实构建了新的计算图，当计算图过多如在迭代中，容易内存/显存溢出。通过 detach 将返回一个新的 Tensor 节点，这个节点孑然一身与其他计算图不相连，即不保留计算图的轨迹。而 detach 是在 requires_grad=False 时才调用的，正因无需计算梯度所以可以只保存为一个单一节点。

for i in range(100):
  sum_loss = sum_loss + x * x
detach function
图1：detach 函数示例：$y=x+1$

Question 1 & 2: Implementing forward and backward computation

在 needle 中，计算图通过 TensorOp 类下属的各个子类进行实际的张量运算，每一个运算符都具有 compute 成员函数，负责前馈计算，本节的多维数组无需自行构建，而是使用 numpy 的 NDArray 实现，为了后续的移植，以 import numpy as array_api 的形式导入 numpy 包。下面对运算符进行依次实现，在本节的代码中将包含每个运算符类的初始化函数（__init__）、前馈函数（compute）以及求导函数（gradient）。

前馈计算直接按照各运算符的定义执行即可，但每个运算符的求导都较为复杂，需要首先明确的是，该求导的对象是损失函数 loss function $\ell$ 对当前节点 $v_i$ 的导数，$v_i$ 是其 inpus 在 对应运算符上的函数，假设 $f=v_i=g(v_{i-1},...v_{i-m})$，即 inputs 为前面的 m 个节点，那么导数为

$$ \frac{\partial \ell}{\partial v_i}=\frac{\partial \ell}{\partial f}\cdot\frac{\partial f}{\partial z_i}\cdot\sum_{j=1}^m\frac{\partial z_i}{\partial z_{i-j}}=\text{out_grad}\cdot\sum_{j=1}^m\frac{\partial z_i}{\partial z_{i-j}} $$

上式是该节点的导数总和，而实际 gradient 返回的是针对每个 inputs 中节点的偏导数，需要注意这一点！

  • PowerScalar：对张量进行乘方运算 $X^n$，例如 $X^2=XX^T$，numpy中提供了相应的乘方接口，直接调用即可

    下面详细推导张量乘方运算的求导函数，假设计算式为 $f=v_2=v_1^n$，则 loss function 对 $v_2$ 的导数可计算为

    $$ \frac{\partial\ell}{\partial v_2}=\frac{\partial\ell}{\partial f}\cdot\frac{\partial f}{\partial v_2}\cdot\frac{\partial v_2}{\partial v_1}=\text{out_grad}\cdot 1\cdot nv_1^{n-1} $$ 
    class PowerScalar(TensorOp):
  def __init__(self, scalar:int):
    self.scalar = scalar

  def compute(self, a: NDArray) -> NDArray:
    return array_api.power(a, self.scalar)

  def gradient(self, out_grad, node):
    a = node.inputs
    return out_grad * self.scalar * array_api.power(a, self.scalar - 1)

  • EWiseDiv：EWise 开头的函数均是逐元素（elementwise）的运算，该函数是对两个张量进行逐元素运算，即 $A \oslash B$

    该函数需要考虑分类讨论三种情况：

    1. 张量 $A$ 和张量 $B$ 的形状（维度）不同，逐元素的运算要求两张量形状相同；
    2. 张量 $B$ 的元素全为 0，虽然在 numpy 中会考虑除 0 的情况，numpy会返回 inf/-inf，符号取决于被除数的符号，在本情况下为 $B$ 的所有元素均为 0；
    3. 正常情况或张量 $B$ 的元素部分为 0，见上一条的解释。

    下面详细推导逐元素除法函数的导函数。由于逐元素，因此可以从单个元素的角度来进行计算，其中会使用类似标量的求导法则，对 $v_3=\frac{v_2}{v_1}$ 如下式所示

    $$ \frac{\partial\ell}{\partial v_3}=\frac{\partial \ell}{\partial f}\cdot\frac{\partial f}{\partial v_3}\cdot(\frac{\partial v_3}{\partial v_2}+\frac{\partial v_3}{\partial v_1})=\text{out_grad}\circ(1\oslash v_1+v_2 \circ (-1) \oslash (v_1\circ v_1)) $$

    其中 $v_2$ 为 lhs (left hand side)，$v_1$ 为 rhs (right hand side)，在下面的运算中均为逐元素的计算，如 numpy 中 * 为逐元素，@ 为矩阵乘法，分别返回对 lhs 和 rhs 的偏导数

    class EWiseDiv(TensorOp):
  def compute(self, a, b):
    if a.shape != b.shape:
      raise RuntimeError("The shapes of the two tensors are inconsistent.")
    if b.all() == 0:
      raise RuntimeError("All elements of the divisor tensor are zero.")
    return array_api.divide(a, b)

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

  • DivScalar：张量逐元素除标量，如 $X/a$，同样需要判断除数是否为 0。numpy 的 divide 函数会自动判别除数的类型，其实在除以标量时，实际上是将标量广播（broadcast）为复制为与被除张量相应维度的矩阵进行逐元素的除法。由于该运算不涉及其他的节点，导函数为

    $$ \frac{\partial \ell}{\partial x}=\frac{\partial\ell}{\partial f}\cdot\frac{\partial f}{\partial x}=\text{out_grad}\cdot\frac{1}{a} $$ 
    class DivScalar(TensorOp):
  def __init__(self, scalar):
    self.scalar = scalar

  def compute(self, a):
    if self.scalar == 0:
      raise RuntimeError("The divisor is zero.")
    return array_api.divide(a, self.scalar)
  def gradient(self, out_grad, node):
    return out_grad / self.scalar

  • MatMul：矩阵乘法 $AB$，numpy 在其中的矩阵乘法函数中已经考虑了两个矩阵的形状的约束。

    对于该函数的导数，假设 $v_3 = v_1 * v_2$，即 lhs 为 $v_1$，rhs 为 $v_2$，那么

    $$ \frac{\partial \ell}{\partial v_3}=\frac{\partial \ell}{\partial f}\cdot\frac{\partial f}{\partial z_3}\cdot(\frac{\partial v_3}{\partial v_1}+\frac{\partial v_3}{\partial v_2}) = \text{out_grad}\cdot(v_2^T+v_1^T) $$

    这里需要考虑形状，假设 $v_3\in\mathbb{R}^{p\times q}, v_1\in\mathbb{R}^{p\times s}, v_2\in\mathbb{R}^{s\times q}$，out_grad 的形状与 $v_3$ 一致，即 $\mathbb{R}^{p\times q}$，则 $\frac{\partial \ell}{\partial v_1}=v_2^T\in\mathbb{R}^{q\times s}$，$\frac{\partial \ell}{\partial v_2}=v_1^T\in\mathbb{R}^{s\times p}$，则 dlhs=matmul(out_grad, rhs.T) 形状为 $\mathbb{R}^{q\times s}$ ，drhs=matmul(lhs.T, outgrad) 形状为 $\mathbb{R}^{s\times p}$，此时符合导数与原函数的维度相同。

    需要注意的是，上述没有考虑在前向传播的过程中使用广播机制，dlhs 的维度可能会比 lhs 的维度多一些，为了使维度统一，需要对这些多余的维度进行求和。例如假设 dlhs.shape = (4, 3, 2)，而 lhs.shape = (3, 2)，那么广播机制在最左侧（通过计算动态得知加入的维度位置）加入了一个维度，需要通过 reduce（实际使用下面的 summation）来减少维度使得维度统一。

    class MatMul(TensorOp):
  def compute(self, a, b):
    return array_api.matmul(a, b)

  def gradient(self, out_grad, node):
    lhs, rhs = node.inputs
    dlhs = array_api.matmul(out_grad, rhs.T)
    drhs = array_api.matmul(lhs.T, out_grad)

    if dlhs.shape != lhs.shape:
      dlhs = summation(dlhs, tuple(range(len(dlhs.shape) - len(lhs.shape))))
    if drhs.shape != rhs.shape:
      drhs = summation(drhs, tuple(range(len(drhs.shape) - len(rhs.shape))))
    assert (dlhs.shape == lhs.shape)
    assert (drhs.shape == rhs.shape)
    return dlhs, drhs

    其中 tuple(range(len(dlhs.shape) - len(lhs.shape))) 动态确定需要求和的维度索引，适用于各种形状的广播情况。例如上面的例子，len(dlhs.shape) - len(lhs.shape) = 1，则 range(1) 为 [0]，tuple 为 (0, )，表示聚合 dlhs 最左侧的维度。通过这样的计算，可以动态获取聚合的轴，从而使得形状对齐。

  • Summation：计算给定轴的和，其中 axes 有多种情况。

    1. axes 为单个整数，即按指定轴求和（0（默认值） 为列，1 为行）；
    2. axes 为元组，如 axes=(0,1)，则沿指定的多个轴求和，在矩阵中此时为求矩阵所有元素的和；
    3. axes 为 None，即求矩阵的所有元素的和
    class Summation(TensorOp):
  def __init__(self, axes: Optional[tuple] = None):
    self.axes = axes

  def compute(self, a):
    return array_api.sum(a, axes=self.axes)

Reference

[1] xx要努力. (2022, 11). Deep Learning System-Homework1. 知乎专栏. https://zhuanlan.zhihu.com/p/579465666.

Contact

There may be some errors present. If you find any, please feel free to contact me at wangqiyao@mail.dlut.edu.cn. I would appreciate it!