https://github.com/tinygrad/teenygrad
简易的深度学习框架，基于CPU上的numpy。在上层封装了Tensor的各种操作，自动微分。是一个用来学习自动微分引擎的比较不错的小项目。代码量也非常小。
同时有一个扩展版本的项目tinygrad，在此之上支持了多种后端，可以看作是pytorch的缩小版

个人感觉需要关注的点主要是三个：

• 数据的表示（Tensor）

• Autograd

• 常见操作的实现

Tensor

Tensor是一个N维的数组，在此之上，支持：

• 多种常见的计算函数

• 自动微分

• 抽象底层数据类型/设备，比如CPU上的数据，GPU上的数据，磁盘上的数据等

Tensor的几个关键变量：

• lazydata，用来抽象底层数据类型。这里取名lazydata的目的是为了做延迟计算。
  • 不过对于TeenyGrad来说，没有做延迟计算，直接认为是一个numpy的NDArray即可

  • Tensor有一个realize的接口，在延迟计算的场景，就是用来做真计算的

• requires_grad，是否需要梯度。因为在训练的时候，只有参数才需要梯度，所以在计算图中的一些旁路分支不需要梯度，就不需要做反向传播了

• grad，保存梯度

• _ctx，用来保存需要做自动微分的上下文。主要是函数的输入值，在反向传播的时候会用到

还有一些关键的函数；

• shape，NDArray每一个维度的大小。
  • 这里还需要意识到一个简化点，就是数据排布不一定是按照维度来的，所以在访问数据的时候，除了shape，一般还需要一个stride的参数，表示每一维增加1的时候，内存中应该移动多少位置。TeenyGrad中是把这个事情也交给numpy维护了
• device，数据的位置。

• dtype，数据的类型，fp32，bf16，long等

Autograd

Autograd的核心点在于通过链式法则来计算梯度，所以需要记录每一个参数到最终的loss的计算过程，这样才能把这个“链”串起来。
计算图就是用来保存这个计算过程的。
举一个Linear的例子：

左图是前向传播，右图是反向传播。从这里就能看出来，计算X/W的梯度的时候，是需要X/W本身的值的。

• 所以在做前向传播的时候，会把一些必要的状态保存下来，在反向传播的时候再用。

• 同时，我们需要记录前向传播时候具体的算子，因为需要根据算子来计算梯度

TeenyGrad中，用Function来表示这些算子，每个算子需要实现两个函数：

• forward，做前向传播，同时把在反向传播阶段需要保存的状态存到Function这个实例中

• backward，用前向传播记录的状态 + 算子本身的梯度计算逻辑，计算对所有参数的梯度。

每一个被计算出来的Tensor，如果需要后续计算梯度的话，就会保留一个_ctx的变量，就是上面说的Function这个实例。

• 同时，每一个算子都有通用的变量parent，用来记录在计算图中的父节点，方便反向传播使用

有了这个ctx，我们只需要遍历计算图，计算每个节点的梯度即可。代码比较简单就直接贴出来了：

  def backward(self) -> Tensor:
    assert self.shape == tuple(), f"backward can only be called for scalar tensors, but it has shape {self.shape})"

    # fill in the first grad with one. don't use Tensor.ones because we don't need contiguous
    # this is "implicit gradient creation"
    self.grad = Tensor(1, device=self.device, requires_grad=False)

    for t0 in reversed(self.deepwalk()):
      assert (t0.grad is not None)
      grads = t0._ctx.backward(t0.grad.lazydata)
      grads = [Tensor(g, device=self.device, requires_grad=False) if g is not None else None
        for g in ([grads] if len(t0._ctx.parents) == 1 else grads)]
      for t, g in zip(t0._ctx.parents, grads):
        if g is not None and t.requires_grad:
          assert g.shape == t.shape, f"grad shape must match tensor shape, {g.shape!r} != {t.shape!r}"
          t.grad = g if t.grad is None else (t.grad + g)
      del t0._ctx
    return self

• 这里reversed(self.deepwalk())是一个逆序的后序遍历，用来做拓扑排序。因为图比较小，所以这样做更方便

• 不断调用_ctx.backward()，使用当前Tensor的梯度，加上保存的ctx，计算出对parent的梯度。然后累加到parent的梯度上。

常见操作

TeenyGrad简化了很多操作，因为并不是针对性能优化的。所以这里尽可能的实现了比较少的算子，总体都在mlops.py这个文件中，一共就200行十来个算子

• 一元的函数，比如Log，Exp等，算算导数就行

• 二元的函数，就是element wise的加减乘除（没有matmul）

• reduction op，比如sum/max，用来消掉某一维

• Movement op，比如Expand，Reshape。expand在反向传播的时候，需要把扩展的那个维度的梯度求和。reshape的话就再reshape回来就行

在tensor.py中基于上面的基本算子实现了很多高级的计算操作。来看几个例子：

matmul

• 上面autograd中，因为matmul太经典了，所以用的matmul的例子。实际上teenygrad没有matmul这个算子

• 比如一个(m, n)乘(n, p)，会先把两个tensor做reshape + tranpose，变成(m, 1, n)和(1, p, n)

• 然后做broadcast + element wise的相乘，变成(m, p, n)，然后对最后一个维度做sum的reduction，得到最终的(m, p)

broadcast

• 这个不是一个具体的算子，而是在执行计算操作的时候都需要的，防止维度没有对齐

• 会先看一下两个Tensor的维度是否相同，如果不相同的话，就会在前面填充(1, )

  • 比如X是(2, 3, 4)，Y是(4, )，那么就会把Y变成(1, 1, 4)。都是一个三维的数组
• 然后针对每一维，进行expand
  • 比如上面的(2, 3, 4), (1, 1, 4)，就会把Y expand成(2, 3, 4)

cumsum

• 做累计和，比如arange就是用这个实现的

• 假设输入[1, 2, 3, 4]

• 会先padding一下，得到[0, 0, 0, 1, 2, 3, 4]

• 然后做pool，得到

  • [0, 0, 0, 1]

  • [0, 0, 1, 2]

  • [0, 1, 2, 3]

  • [1, 2, 3, 4]

• 然后针对pool出来的矩阵做sum，就得到了[1, 3, 6, 10]

• 这里实现还做了一个分块。如果数据量比较大的时候，会先分一下块，针对每一个块做cumsum，然后最后再做一下累加

其实内部还有很多复杂的操作，这里就不多细抠了，感兴趣可以自己去看一看这些实现的方法，当作口香糖嚼一嚼

• 比如上面看到的pool，需要涉及到大量的形状变化。

• 比如py的__getitem__，可以传入一个tuple，取多个维度的切片。实现起来也是非常复杂。同时这些也是要参与到反向传播的，所以需要使用上面这些特定的算子。

• 相比之下一些和形状关联不大的计算操作就比较简单。

除此之外，TeenyGrad还提供了几个optimizer的简单实现，和算法里描述的一样，就不单独列了。