动机、参考资料、涉及内容

直接动机：

  • CTC 是怎么计算梯度的【大概最终不涉及】

其他动机：

  • 按 CS231N 中描述的，反向传播算法里存在由前向后与由后向前两种方式，与 tape 这个东西有关，tensorflow的API就明确有 tf.GradientTape 这个东西，因此有必要研究清楚
  • micrograd 与 tinygrad 这两个 Github 项目有实现一套自动微分框架，需要研究一下
  • 卷积操作的反向传播也是卷积操作？需要仔细推导一次【这个在陈天奇的课程中有介绍】

参考资料

本文的内容主要是对下面参考资料的“注解”

链式法则【貌似没啥用】

$n$ 个输入, $m$ 个输出, 定义 Jacobean 矩阵为 (形状为 $m\times n$)

\[J = \begin{bmatrix} \frac{\partial y_1}{\partial x_1} & \cdots & \frac{\partial y_1}{\partial x_n} \\ \vdots & \ddots &\vdots \\ \frac{\partial y_m}{\partial x_1} & \cdots & \frac{\partial y_m}{\partial x_n} \end{bmatrix}= \begin{bmatrix} \nabla^T {y_1} \\ \vdots \\ \nabla^T {y_m} \\ \end{bmatrix}\]

假设: $\mathbf{z}=f(\mathbf{y}), y=g(\mathbf{x})$, 即: $\mathbf{z}=f(g(\mathbf{x}))$, 记 $\nabla\mathbf{x}=J(\mathbf{x}, \mathbf{z})$, 链式法则可以写为

\[\nabla \mathbf{x}= \begin{bmatrix} \frac{\partial L}{\partial x_1} \\ \vdots \\ \frac{\partial L}{\partial x_n} \\ \end{bmatrix} = J^T \begin{bmatrix} \frac{\partial L}{\partial x_1} \\ \vdots \\ \frac{\partial L}{\partial x_n} \end{bmatrix} =J^T(\nabla \mathbf{y})\]

Sum-Product 算法

对于一个有向无环图, 我们把只有出边的节点称为叶子节点, 只有入边的节点称为根节点, 假设存在一条边 $v_i\to v_j$, 那么称 $v_j$ 是 $v_i$ 的父节点. 一般来说, 一个有向无环图中可以有多个叶子节点和多个根节点, 我们这里只讨论只有一个根节点的情形.

备注:

  • 注意一些概念上的区分, 【树】是一种特殊的【只有一个根节点的有向无环图】
  • 【只有一个根节点的有向无环图】有一个性质是: 所有其他节点到根节点都至少存在一条路

我们在只有一个根节点的有向无环图上考虑一个这样的问题:

  • 边的值: 假设每条边都有一个值, 例如边 $v_i\to v_j$ 上的值记作 $s_{i,j}$
  • 根节点的值: 我们定义【根节点的值】为 $1$
  • 到根节点的路的值: 【一个节点到根节点的一条路】的值被定义为其经过的所有边的值的乘积
  • 其他节点的值: 【一个节点的值】定义为【其所有到根节点的路】的值之和

现在我们给定每条边的值, 需要按上面的定义计算每个节点的值. 后面我们将看到, 自动微分算法的本质就是这个过程.

自动微分原理

本节结合 陈天奇课程 的第 4 节课和 Autodiff 博客 做原理上的讲解

实现方案比较及阅读建议

  • Autodiff: 支持高阶导数以及向量化(简易版), 可以说是最简实现(代码行数), 实现上没有显式使用拓扑排序, 而是采用了递归, 可能需要仔细揣摩.
  • Micrograd: 不支持高阶导数, 支持向量化(简易版), 相当于 Autodiff 的简化版本, 实现上显式使用了拓扑排序, 便于理解.
  • dlsyscourse: 完整框架, 支持高阶导数, 支持向量化(高效实现), 课程内容非常丰富, 适合认真学习, 获益很多.
  • Simple Autograd: 支持高阶导数, 支持向量化, 还考虑了一些诸如自定义算子时的一些问题(初看时想了很久不得要领), 这部分内容不确定 dlsyscourse 里怎么涉及的
  • Tinygrad: 待研究
  • CS231N: 完成前面的内容后可以审视一下, 检验理解是否正确

Python 高级特性【待定】

要完全理解具体的实现, 实际上需要对一些 Python 高级特性进行理解

  • closure
  • 内置方法 sum 的执行逻辑
  • 运算符重载: __add__, __radd__
  • __new__ 与 __init__: 目前感觉只是 dlsyscourse 里故弄玄虚?

__add__ 与 __radd__

Python 官方文档: https://docs.python.org/3/reference/datamodel.html#emulating-numeric-types

当遇到表达式 x + y 时:

  • 首先尝试使用 type(x).__add__(x, y), 如果 x 没有 __add__ 方法 (或者返回值是 NotImplemented)
  • 则再尝试调用 type(y).__radd__(y, x)

当遇到表达式 x += y 时

  • 首先尝试调用 x = type(x).__iadd__(x, y)
  • 然后再尝试 x = type(x).__add__(x, y)
  • 最后尝试 x = type(y).__radd__(y, x)

各种运算符重载可参见上面的官方文档

class A:
    def __init__(self, value):
        self.value = value
    def __add__(self, other):    # 次优先
        return A(self.value + other.value)
    def __iadd__(self, other):   # 最优先
        return 7
    def __mul__(self, other):
        return NotImplemented    # NotImplemented 是一个 Python 定义的常量, 类似于 True/False

class B:
    def __init__(self, value):
        self.value = value
    def __radd__(self, other):   # 最次
        return B(self.value + other.value)

a = A(1)
b = B(2)
a + b    # A(3), B(3)
a += b   # 7, A(3), B(3)

注意: 一般不能将 __rsub__ 与 __sub__ 定义成一样:

class C:
    def __init__(self, value):
        self.value = value
    def __sub__(self, other):
        print("__sub__ called")
        if isinstance(other, C):
            return C(self.value - other.value)
        else:
            return C(self.value - other)
    def __rsub__(self, other):
        print("__rsub__ called")
        if isinstance(other, C):
            return C(other.value - self.value)
        else:
            return C(other - self.value)
    def __str__(self):
        return f"C({self.data})"
    def __repr__(self):
        return f"C({self.data})"

print(1 - C(5))     # __rsub__: C(-4)
print(C(1) - 5)     # __sub__:  C(-4)
print(C(1) - C(5))  # __sub__:  C(-4)

sum 与 np.sum

python 内置的 sum 函数实际上有一个默认参数 start, 其默认值为 0, 因此执行 sum([x, y]) 时实际上会触发两次加法, 即: 0 + x + y.

sum([1, 2], start=6)  # 结果是 9, 等价于: 6 + 1 + 2, 因此会触发 2 次假发

而 np.sum 没有类似的 start 参数, 因此会少触发一次加法

class Data:
    def __init__(self, data):
        self.data = data
    def __add__(self, other):
        print(f"__add__({self}, {other}) called")
        if isinstance(other, Data):
            return Data(self.data + other.data)
        else:
            return Data(self.data + other)
    def __radd__(self, other):
        print(f"__radd__({self}, {other}) called")
        return self.__add__(other)
    def __str__(self):
        return f"Data({self.data})"
    def __repr__(self):
        return f"Data({self.data})"

# convert NumPy array into array of Data objects:
to_data = np.vectorize(lambda x : Data(x))

# get values from array of Data objects:
to_vals = np.vectorize(lambda data : data.data)

x = list(range(3))

np.sum(to_data(x))  # 返回结果: Data(3)
# 输出结果:
# __add__(Data(0), Data(1)) called
# __add__(Data(1), Data(2)) called

sum(to_data(x))             # 返回结果: Data(3)
# 输出结果
# __radd__(Data(0), 0) called
# __add__(Data(0), 0) called
# __add__(Data(0), Data(1)) called
# __add__(Data(1), Data(2)) called

Autodiff

这里直接搬运了博客中的实现如下

自动微分的实现(只能计算一阶导数)

from collections import defaultdict

class Variable:
    def __init__(self, value, local_gradients=()):
        self.value = value
        self.local_gradients = local_gradients
    
def add(a, b):
    "Create the variable that results from adding two variables."
    value = a.value + b.value    
    local_gradients = (
        (a, 1),  # the local derivative with respect to a is 1
        (b, 1)   # the local derivative with respect to b is 1
    )
    return Variable(value, local_gradients)

def mul(a, b):
    "Create the variable that results from multiplying two variables."
    value = a.value * b.value
    local_gradients = (
        (a, b.value), # the local derivative with respect to a is b.value
        (b, a.value)  # the local derivative with respect to b is a.value
    )
    return Variable(value, local_gradients)

def get_gradients(variable):
    """ Compute the first derivatives of `variable` 
    with respect to child variables.
    """
    gradients = defaultdict(lambda: 0)
    
    def compute_gradients(variable, path_value):
        for child_variable, local_gradient in variable.local_gradients:
            # "Multiply the edges of a path":
            value_of_path_to_child = path_value * local_gradient
            # "Add together the different paths":
            gradients[child_variable] += value_of_path_to_child
            # recurse through graph:
            compute_gradients(child_variable, value_of_path_to_child)
    
    compute_gradients(variable, path_value=1)
    # (path_value=1 is from `variable` differentiated w.r.t. itself)
    return gradients

使用:

a = Variable(4)
b = Variable(3)
c = add(a, b) # = 4 + 3 = 7
d = mul(a, c) # = 4 * 7 = 28

gradients = get_gradients(d)

print('d.value =', d.value)
print("The partial derivative of d with respect to a =", gradients[a])
print('gradients[b] =', gradients[b])
print('gradients[c] =', gradients[c])
print('dict(d.local_gradients)[a] =', dict(d.local_gradients)[a])
print('dict(d.local_gradients)[c] =', dict(d.local_gradients)[c])
print('dict(c.local_gradients)[a] =', dict(c.local_gradients)[a])
print('dict(c.local_gradients)[b] =', dict(c.local_gradients)[b])

功能说明:

  • 这个实现只能支持标量计算(应该可以比较容易扩展到张量), 且只能计算一阶导数
  • compute_gradients 是一个递归实现, 入参 path_value 实际上是从跟节点(Variable)到当前节点(variable)的一条路径的局部梯度的乘积, compute_gradients 的是将这条路径的所有后继梯度乘积进行累积

注意: 不能随意进行 inplace 操作

class Variable:
    def __init__(self, value, local_gradients=()):
        self.value = value
        self.local_gradients = local_gradients
    def __mul__(self, other):
        return mul(self, other)
    def __str__(self):
        return f"Variable({self.value})"
    def __repr__(self):
        return f"Variable({self.value})"

def mul(a, b):
    value = a.value * b.value
    local_gradients = (
        (a, lambda path_value: path_value * b),
        # local gradient for a is b, so multiply path_value by b.
        (b, lambda path_value : path_value * a)
        # local gradient for b is a, so multiply path_value by a.
    )
    return Variable(value, local_gradients)

a, b = Variable(2), Variable(3)
c = mul(a, b)
b.value = 10   # 这个相当于是inplace的操作, 可能会引发问题
c.local_gradients[0][1](Variable(1))  # Variable(10)

Micrograd

Needle (陈天奇课程记录与实现)【待定】

pytorch 自动微分的一些探索

本节是临时记录, 主要目的是为 Simple Autograd 的一些内容做补充

关于 in-place 操作

a = torch.tensor(1., requires_grad=True)
b = torch.tensor(2., requires_grad=True)
# a.random_()  # Error: 叶子节点不能用 inplace 操作
c = a * b
d = a + b
c.random_()    # 中间节点用 inplace 操作, 猜测实现上应该是把 c 的梯度传递变成了 0
e = c + d
e.backward()
print(a.grad, b.grad)  # 1, 1

关于 grad_fn

a = torch.tensor(1., requires_grad=True)
b = torch.tensor(2., requires_grad=True)
c = a * b
c.grad_fn(torch.tensor(1.))  # 手动调用 grad_fn: (tensor(2., grad_fn=<MulBackward0>), tensor(1., grad_fn=<MulBackward0>))
c.random_()
c.grad_fn(torch.tensor(1.))  # tensor(0.)

注意第一个结果里返回的结果里 tensor 是有 grad_fn 的, 而第二个是没有的

高阶导数

torch.random.manual_seed(42)
x = torch.rand((4,), requires_grad=True)
y = torch.rand((4,), requires_grad=True)
print(x, y)

# 使用 torch.autograd.grad 接口
z = (x * y).sum()
dz_dx, dz_dy = torch.autograd.grad(z, [x, y], create_graph=True)
w = (dz_dx * dz_dy).sum()
dw_dx, dw_dy = torch.autograd.grad(w, [x, y], create_graph=True)
print(dw_dx, dw_dy)  # (x.grad, y.grad) = (None, None)

# 使用 torch.Tensor.backward 接口
z = (x * y).sum()
dz_dx, dz_dy = torch.autograd.grad(z, [x, y], create_graph=True)
w = (dz_dx * dz_dy).sum()
w.backward()
print(x.grad, y.grad)

通常情况下, 自定义算子不能求高阶导数?

# 备注: 这里是按 torch 1.12 版本的写法来写的, 即没有写独立的 setup_context 函数
class LinearFunction(Function):
    @staticmethod
    def forward(ctx, input, weight, bias):
        ctx.save_for_backward(input, weight, bias)
        output = input.mm(weight.t())
        if bias is not None:
            output += bias.unsqueeze(0).expand_as(output)
        return output
    @staticmethod
    def backward(ctx, grad_output):
        input, weight, bias = ctx.saved_tensors
        grad_input = grad_weight = grad_bias = None
        if ctx.needs_input_grad[0]:
            grad_input = grad_output.mm(weight)
        if ctx.needs_input_grad[1]:
            grad_weight = grad_output.t().mm(input)
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum(0)

        return grad_input, grad_weight, grad_bias

linear = LinearFunction.apply

x, w, b = (torch.rand((2, 3)), torch.rand((4, 3), requires_grad=True), torch.rand((4,), requires_grad=True))
z = linear(x, w, b).sum()
dz_dw, dz_db = torch.autograd.grad(z, [w, b], create_graph=True)
print(dz_dw, dz_db)   # 只是 Tensor, 没有 grad_fn
y = (dz_dw + dz_db.reshape(-1, 1)).sum()
dy_dw, dy_db = torch.autograd.grad(y, [w, b], create_graph=True)  # Error

Simple Autograd 注解

本节是对 simple autograd 这个 colab notebook 的一个注解

torch/csrc/jit/runtime/autodiff.cpp

对于这个 notebook 中作为示例的高阶导数, 这里先手动求出结果, 以便后续分析:

\[\mathbf{a}=(a_1,\ldots,a_n), \mathbf{b}=(b_1,\ldots,b_n)\\ L_1=\sum_{i=1}^{n}(a_i+b_i)b_i\\ \frac{\partial L_1}{\partial \mathbf{a} }=\mathbf{b}, \frac{\partial L_1}{\partial \mathbf{b} }=\mathbf{a} + 2\mathbf{b}\\ L_2=\text{sum}[(\frac{\partial L_1}{\partial \mathbf{a} })^2+(\frac{\partial L_1}{\partial \mathbf{b} })^2]=\sum_{i=1}^{n}{[b_i^2+(a_i+2b_i)^2]} \\ \frac{\partial L_2}{\partial \mathbf{a} }=2\mathbf{a}+4\mathbf{b}, \frac{\partial L_2}{\partial \mathbf{b} }=4\mathbf{a}+10\mathbf{b}\]

Tinygrad (早期版本)

在 tinygrad 最初的一些版本里, 例如: 17bf90db, tinygrad 宣称总代码行数将不超过 1000 行, 适合本文进行基本的原理研究

Tinygrad (近期版本)

近期 tinygrad 逐渐朝着一个真正的深度学习框架进行发展

git log -S"tinygrad will always be below 1000 lines" -- README.md

# 以下是输出内容
# commit 803587b8b42697fcc2f5ad920ea76d9cd716f666
# Author: George Hotz <geohot@gmail.com>
# Date:   Fri May 26 06:10:41 2023 +0000

#     update readme

# commit 2e7f16bf3f9b7b2934073acd3bae7f864c3b6ddf
# Author: George Hotz <geohot@gmail.com>
# Date:   Mon Nov 2 07:42:11 2020 -0800

#     the power of cheating

git show 803587b8b42697fcc2f5ad920ea76d9cd716f666

CS231N【待定】

留作 Code Review