手写深度学习框架(一)——基础数据结构与线性层

本文旨在说明如何使用 Python 和 NumPy 从零开始，实现一个基础的自动求导引擎。这个引擎的核心功能是构建计算图并执行反向传播，特别是对神经网络中常见的矩阵运算提供支持。这一部分的内容是用自建的引擎定义一个包含线性层的简单神经网络，并通过训练使其拟合一个给定的函数，以此来验证其功能的正确性。主要参考了《深度学习入门2自制框架 (斋藤康毅)》其中的架构。完整代码在文末。

代码主要由以下几个部分构成：

• Value 类 : 核心数据结构，用于封装 numpy 数组，并追踪其在计算图中的梯度和来源。
• Function 类 : 所有运算的基类，定义了前向传播和反向传播的接口。
• Module 与 Linear 类 : 模仿常见深度学习框架的模块化设计，用于构建神经网络层。
• 训练示例 : 使用上述组件来实际训练一个模型。

1. Value 类：计算图的核心节点

为了实现自动求导，需要将所有的计算过程表示为一个计算图。Value 类就是这个图中的节点，它包装了 numpy 数组，并额外存储了与梯度计算相关的信息。

计算图相关可以看在B站上面查看相关视频

Value 类的主要属性如下：

• data: 存储实际的数值，类型为 numpy.ndarray，以便支持向量和矩阵。
• grad: 存储该节点相对于最终输出的梯度。在反向传播过程中，这个值会被计算并填充。
• created_by: 指向一个 Function 对象，记录该 Value 是由哪个运算产生的。这形成了节点之间的连接，构成了计算图。
• generation: 一个整数，表示节点在计算图中的“深度”。输入数据为第 0 代，每经过一次运算，代数加一。这个属性用于在反向传播时进行拓扑排序。

import numpy as np

class Value:
    """一个存储标量或矩阵并记录其梯度的容器。"""
    def __init__(self, data, name=""):
        if not isinstance(data, np.ndarray):
            self.data = np.array(data)
        else:
            self.data = data

        self.grad = np.zeros_like(self.data)  # 梯度初始化为0
        self.created_by: Function | None = None  # 记录节点的创建者
        self.generation = 0
        self.name = name

    def set_creator(self, func):
        self.created_by = func
        self.generation = func.generation + 1

    # backward 方法将在后面介绍
    # ...

2. Function 类：定义运算

Function 类是计算图中的“边”，代表一个具体的运算，如加法或矩阵乘法。它定义了所有运算的通用结构。

• forward(*xs) : 执行具体的计算。输入都转化为 numpy 数组，输出为 Value，默认只有一个输出。
• backward(*gys) : 计算梯度。根据输出端的梯度 gys（来自链式法则的上一环），计算并返回输入端的梯度 gxs。

通过 __call__ 方法，我们可以像普通函数一样调用一个 Function 实例。这个方法会完成从 Value 对象中提取 data，执行 forward 计算，然后将结果包装回一个新的 Value 对象，并设置好 created_by 关系，从而将新节点连接到计算图中。

class Function:
    """
    运算的基类。它连接输入的Value对象和输出的Value对象。
    """

    def __init__(self):
        self.inputs: list[Value] = []  # 存储函数的输入
        self.outputs: list[Value] = []  # 存储函数的输出
        self.generation: int = -1  # 函数的代数

    def forward(self, *xs: np.ndarray) -> Value:
        raise NotImplementedError()

    def backward(self, gy: np.ndarray) -> tuple[np.ndarray]:
        raise NotImplementedError()

    def __call__(self, *inputs) -> Value:
        """使得类可以想普通函数一样调用"""
        # 首先先对输入进行转换，确保数据格式一致
        xs = [as_array(x) for x in inputs]
        self.inputs = [as_value(x) for x in inputs]
        y = self.forward(*xs)  # 调用 forward 方法
        y = as_value(y)  # 转换为 Value 对象
        self.outputs = [y]
        y.set_creator(self)  # 设置创建者
        return y

加法Add实现

加法的反向传播就将输入梯度返回给两个数即可，这里需要处理梯度广播，我们线性层的输入是(batch_size,in_feature)，而线性层中的偏置bias是(in_feature,)，因此我们到bias上的梯度是(batch_size,in_feature)，需要将其按行相加，然后设置为bias的梯度。

class Add(Function):
    """加法运算"""

    def forward(self, x, y):
        return x + y

    def backward(self, gy) -> tuple[np.ndarray, np.ndarray]:
        # 加法的梯度会原封不动地分配给两个输入
        x, y = self.inputs
        if y.data.shape != gy.shape:  # 如果输入的形状不一致，则进行广播
            gy_bias = np.sum(gy, axis=0)
        return gy, gy_bias

def add(x, y):
    return Add()(x, y)

矩阵乘法的实现

矩阵乘法是神经网络中的一个核心运算。其前向传播为\(Y = X * W.T\)。根据矩阵求导法则，其反向传播（梯度计算）如下：

\(\frac{\partial L}{\partial X} = \frac{\partial L}{\partial Y} \cdot W^T\)

\(\frac{\partial L}{\partial W} = X^T \cdot \frac{\partial L}{\partial Y}\)

代码实现直接对应上述公式：

class MatMul(Function):
    """矩阵乘法运算"""
    def forward(self, x, y):
        # 保存输入，用于反向传播
        self.matmul_inputs = [x, y]
        return np.dot(x, y)

    def backward(self, grad_output: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        x, y = self.matmul_inputs
        grad_x = np.dot(grad_output, y.T)
        grad_y = np.dot(x.T, grad_output)
        return grad_x, grad_y

def matmul(x, y):
    return MatMul()(x, y)

矩阵转置的实现

class Transpose(Function):
    """转置"""

    def forward(self, x):
        self.input = x
        return x.T

    def backward(self, grad_output):
        return grad_output.T


def transpose(x):
    f = Transpose()
    return f(x)

Relu的实现

class Relu(Function):
    def forward(self, x):
        return np.maximum(x, 0.0)

    def backward(self, gy):
        x = self.inputs[0].data
        mask = x > 0
        return gy * mask


def relu(x):
    return Relu()(x)

3. backward 方法：梯度计算引擎

backward 方法是自动求导功能的核心，它被定义在 Value 类中，当在最终输出节点上调用时，会触发整个计算图的梯度回传。

其工作流程如下：

1. 初始化与排序 : 从调用 backward 的 Value 节点开始，反向遍历计算图。利用 generation 属性对所有相关的 Function 进行拓扑排序，确保从“后代”节点往“祖先”节点方向计算。
2. 梯度计算 : 按照排序，依次处理队列中的每一个 Function。首先获取该 Function 输出节点的梯度（这个梯度是之前步骤中计算好的），然后调用该 Function 自身的 backward 方法。
3. 梯度累积 : backward 方法返回对其输入的梯度。这些梯度值会被累加到对应输入 Value 对象的 .grad 属性上。使用累加（+=）是因为一个节点可能作为多个后续运算的输入，其总梯度是所有路径传回的梯度之和。
4. 迭代 : 将当前处理过的节点的父节点也加入待处理队列中，重复此过程，直到遍历完所有相关的节点。

Python

# 在 Value 类中
def backward(self):
    """执行反向传播，计算图中所有节点的梯度。"""
    funcs: list[Function] = []
    visited = set()

    # 最终节点的梯度默认为1
    if self.grad is None:
        self.grad = np.ones_like(self.data)

    def add_func(func: Function):
        if func not in visited:
            funcs.append(func)
            visited.add(func)
            funcs.sort(key=lambda x: x.generation)

    assert self.created_by is not None, "不能在输入节点上调用反向传播"
    add_func(self.created_by)

    while funcs:
        f = funcs.pop() # 取出 generation 最高的函数
        gys = [output.grad for output in f.outputs]
        gxs = f.backward(*gys)

        for x, gx in zip(f.inputs, gxs):
            # 累积梯度
            if x.grad is None:
                x.grad = gx
            else:
                x.grad = x.grad + gx

            if x.created_by is not None:
                add_func(x.created_by)

4. Module 和 Linear：构建神经网络层

为了代码结构化和参数管理，我们引入 Module 基类。它的主要作用是递归地收集一个模块及其子模块中所有可学习的参数（即 Value 对象）。

class Module:
    """神经网络模块的基类"""
    def parameters(self):
        params = []
        for name, value in self.__dict__.items():
            if isinstance(value, Value):
                params.append(value)
            elif isinstance(value, Module):
                params.extend(value.parameters())
        return params

    def zero_grad(self):
        """将所有参数的梯度清零"""
        for p in self.parameters():
            p.grad = np.zeros_like(p.grad)

Linear 层（全连接层）是 Module 的一个具体实现。权重的维度采用(out_feature, in_feature)

class Linear(Module):
    """线性层，实现 y = x @ W.T + b"""
    def __init__(self, in_features, out_features, bias=True, name=""):
        # ... 初始化 ...
        k = np.sqrt(2.0 / in_features)
        self.weight = Value(np.random.uniform(-k, k, (out_features, in_features)), name=name + "_weight")
        if bias:
            self.bias = Value(np.zeros(out_features), name=name + "_bias")

    def __call__(self, x):
        a = x @ self.weight.T
        return a + self.bias

权重 weight 和偏置 bias 被定义为 Value 对象，因此它们的梯度可以被自动计算。

5. 训练示例

最后，我们用以上组件构建一个三层神经网络，来拟合函数\(y = x^2 + 2x + 1\)。
训练过程遵循标准的流程：

1. 前向传播 : 输入数据 X_train 通过模型，计算得到预测值 Y_pred。
2. 计算损失 : 用均方误差（MSE）计算 Y_pred 和真实值 Y_train 的差距。
3. 反向传播 : 首先手动计算损失对 Y_pred 的初始梯度。对于 MSE，该梯度为 frac2N(y_pred−y_true)。然后调用 Y_pred.backward()，自动计算出模型中所有参数的梯度。
4. 参数更新 : 使用梯度下降法更新模型的所有参数：param.data -= learning_rate * param.grad。
5. 梯度清零 : 为下一次迭代做准备。

# 定义模型
model = MyModel() # MyModel 由 l1, l2, l3 三个 Linear 层构成

# 训练循环
for epoch in range(epochs):
    # 1. 前向传播
    Y_pred = model(Value(X_train))

    # 2. 计算损失
    loss_val = np.mean((Y_pred.data - Y_train) ** 2)

    # 3. 反向传播
    grad_initial = 2 * (Y_pred.data - Y_train) / len(X_train)
    Y_pred.grad = grad_initial
    Y_pred.backward()

    # 4. 更新参数
    for p in model.parameters():
        p.data -= learning_rate * p.grad

    # 5. 清零梯度
    model.zero_grad()

运行该训练脚本，可以看到损失值逐渐降低，模型最终可以对给定的输入做出较为准确的预测，证明了我们实现的自动求导引擎是有效的。

完整代码

import numpy as np


class Value:
    """
    一个存储单个标量并记录其梯度的容器。
    它构建了一个计算图，用于自动求导。
    """

    def __init__(self, data, name=""):
        # 确保所有数据都是 numpy 数组，方便后续运算
        if not isinstance(data, np.ndarray):
            # 将标量或列表转换为 numpy 数组
            self.data = np.array(data)
        else:
            self.data = data

        self.grad = np.zeros_like(self.data)  # 梯度初始化为0
        self.created_by: Function | None = None  # 这个value是通过函数创建的，如果是输入，则为None
        self.generation = 0
        self.name = name

    def __repr__(self):
        return f"Value(data={self.data})"

    def set_creator(self, func):
        """
        设置这个value的创建者
        """
        self.created_by = func
        self.generation = func.generation + 1

    @property
    def T(self):
        return transpose(self)

    def backward(self):
        """
        执行反向传播，计算图中所有节点的梯度。
        """
        funcs: list[Function] = []  # 创建一个空列表，用于存储所有节点的函数
        visited = set()
        if self.grad is None:
            self.grad = np.ones_like(self.data)

        def add_func(func: Function):
            """添加一个函数"""
            if func not in visited:
                funcs.append(func)
                visited.add(func)
            funcs.sort(key=lambda x: x.generation)

        assert self.created_by is not None, "无法在非函数生成变量中调用反向传播"
        add_func(self.created_by)  # 添加创建该节点的函数

        while funcs:
            f = funcs.pop()  # 取出代数最大的函数
            gys = [output.grad for output in f.outputs]  # 获取输出节点的梯度
            gxs = f.backward(*gys)
            for x, gx in zip(f.inputs, gxs):
                if x.grad is None:  # 如果该节点的梯度没有被计算过
                    x.grad = gx
                else:  # 如果该节点的梯度已经被计算过
                    x.grad = x.grad + gx
                if x.created_by is not None:
                    add_func(x.created_by)  # 添加该节点的creator

    def __add__(self, other):
        return add(self, other)

    def __radd__(self, other):
        return add(other, self)

    def __matmul__(self, other):
        return matmul(self, other)


def as_array(x: Value | np.ndarray | float | int | None) -> np.ndarray:
    """将各种类型转换为numpy数组"""
    if x is None:
        raise ValueError("输入不能为None")
    if isinstance(x, Value):
        return x.data
    if isinstance(x, np.ndarray):
        return x
    if np.isscalar(x):
        return np.array(x)
    raise TypeError(f"不支持的输入类型: {type(x)}")


def as_value(x: Value | np.ndarray | float | int | None) -> Value:
    """将各种类型转换为Value"""
    if x is None:
        raise ValueError("输入不能为None")
    if isinstance(x, Value):
        return x
    if isinstance(x, np.ndarray):
        return Value(x)
    if np.isscalar(x):
        return Value(np.array(x))
    raise TypeError(f"无法将{type(x)}转换为Value")


# -------------------------------------------
# 2. Function 类及其子类: 实现具体的运算功能
# -------------------------------------------
class Function:
    """
    运算的基类。它连接输入的Value对象和输出的Value对象。
    """

    def __init__(self):
        self.inputs: list[Value] = []  # 存储函数的输入
        self.outputs: list[Value] = []  # 存储函数的输出
        self.generation: int = -1  # 函数的代数

    def forward(self, *xs: np.ndarray) -> Value:
        raise NotImplementedError()

    def backward(self, gy: np.ndarray) -> tuple[np.ndarray]:
        raise NotImplementedError()

    def __call__(self, *inputs) -> Value:
        """使得类可以想普通函数一样调用"""
        # 首先先对输入进行转换，确保数据格式一致
        xs = [as_array(x) for x in inputs]
        self.inputs = [as_value(x) for x in inputs]
        y = self.forward(*xs)  # 调用 forward 方法
        y = as_value(y)  # 转换为 Value 对象
        self.outputs = [y]
        y.set_creator(self)  # 设置创建者
        return y


class Add(Function):
    """加法运算"""

    def forward(self, x, y):
        return x + y

    def backward(self, gy) -> tuple[np.ndarray, np.ndarray]:
        # 加法的梯度会原封不动地分配给两个输入
        x, y = self.inputs
        if y.data.shape != gy.shape:  # 如果输入的形状不一致，则进行广播
            gy_bias = np.sum(gy, axis=0)
        return gy, gy_bias


def add(x, y):
    return Add()(x, y)


class MatMul(Function):
    """矩阵乘法运算"""

    def forward(self, x, y):
        # 在前向传播时，保存输入用于反向传播
        self.matmul_inputs = [x, y]
        return np.dot(x, y)

    def backward(self, grad_output: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        x, y = self.matmul_inputs
        # 根据矩阵乘法的求导法则计算梯度
        # (A @ B)' = A' @ B^T + A^T @ B'
        assert isinstance(x, np.ndarray) and isinstance(y, np.ndarray), "输入必须是 numpy 数组"
        return np.dot(grad_output, y.T), np.dot(x.T, grad_output)


class Relu(Function):
    def forward(self, x):
        return np.maximum(x, 0.0)

    def backward(self, gy):
        x = self.inputs[0].data
        mask = x > 0
        return gy * mask


def relu(x):
    return Relu()(x)


def matmul(x, y):
    f = MatMul()
    return f(x, y)


class Transpose(Function):
    """转置"""

    def forward(self, x):
        self.input = x
        return x.T

    def backward(self, grad_output):
        return grad_output.T


def transpose(x):
    f = Transpose()
    return f(x)


# ---------------------------------
# 3. 模块/层 的实现
# ---------------------------------
class Module:
    """所有神经网络模块的基类"""

    def parameters(self):
        """返回模块中所有可学习的参数 (Value对象)"""
        params = []
        for name, value in self.__dict__.items():
            if isinstance(value, Value):
                params.append(value)
            elif isinstance(value, Module):
                params.extend(value.parameters())
        return params

    def zero_grad(self):
        """将所有参数的梯度清零"""
        for p in self.parameters():
            p.grad = np.zeros_like(p.grad)


class Linear(Module):
    """
    线性层，实现 y = x @ W.T + b
    """

    def __init__(self, in_features, out_features, bias=True, name=""):
        # 输入是(batch_size, in_features)，输出是(batch_size, out_features)
        self.in_features = in_features
        self.out_features = out_features

        # He/Kaiming 初始化
        # 修改 Linear.__init__()
        k = np.sqrt(2.0 / in_features)  # He 初始化
        self.weight = Value(np.random.uniform(-k, k, (out_features, in_features)), name=name + "_weight")

        self.use_bias = bias
        if self.use_bias:
            self.bias = Value(np.zeros(out_features), name=name + "_bias")  # 列向量

    def __call__(self, x):
        a = x @ self.weight.T  # 得到的是(batch_size, out_features)
        a.name = "weight @ x"
        b = a + self.bias
        b.name = "weight @ x + bias"
        return b


class MyModel:
    def __init__(self):
        self.l1 = Linear(1, 4, name="l1")
        self.l2 = Linear(4, 4, name="l2")
        self.l3 = Linear(4, 1, name="l2")
        self.input: Value | None = None
        self.output: Value | None = None
        self.params = self.l1.parameters() + self.l2.parameters()

    def __call__(self, x):
        return self.forward(x)

    def forward(self, x):
        x = self.l1(x)
        x = relu(x)
        x = self.l2(x)
        x = relu(x)
        x = self.l3(x)
        self.output = x
        return x

    def zero_grad(self):
        self.l1.zero_grad()
        self.l2.zero_grad()
        self.l3.zero_grad()

    def parameters(self):
        """返回模块中所有可学习的参数 (Value对象)"""
        params = []
        params = self.l1.parameters() + self.l2.parameters() + self.l3.parameters()
        return params


# ---------------------------------
# 4. 测试函数: 拟合 y = 2x
# ---------------------------------
def test_fit():
    print("开始测试拟合函数 y = 2x ...")

    # a. 创建数据集
    # 输入维度为1，输出维度为1
    def generate_data(start, end, num_points, f):
        """
        根据线性方程 y = 2x 生成数据点。

        参数:
        start (float): 输入范围的起始值
        end (float): 输入范围的结束值
        num_points (int): 要生成的数据点数量

        返回:
        tuple: 包含生成的输入特征和输出标签的元组
        """
        X = np.linspace(start, end, num_points).reshape(-1, 1)
        Y = f(X)
        return X, Y

    def f(x):
        return x**2 + 2 * x + 1

    X_train, Y_train = generate_data(1, 10, 100, lambda x: f(x))

    # b. 定义模型
    # 我们的模型就是一个简单的线性层
    # 输入特征是1维，输出特征也是1维
    model = MyModel()

    # c. 定义超参数和优化器
    learning_rate = 0.01
    epochs = 10000

    # d. 训练循环
    for epoch in range(epochs):
        # 前向传播
        Y_pred = model(Value(X_train))
        # print(Y_pred.data.shape)

        # 计算损失 (均方误差 MSE)
        loss_val = np.mean((Y_pred.data - Y_train) ** 2)

        # 手动设置损失对预测值的初始梯度
        grad_initial = 2 * (Y_pred.data - Y_train) / len(X_train)
        Y_pred.grad = grad_initial

        # 反向传播
        Y_pred.backward()

        # 更新参数 (SGD)
        for p in model.parameters():
            p.data -= learning_rate * p.grad

        if epoch % 100 == 0 or epoch == epochs - 1:
            # 权重应该接近 2，偏置应该接近 0
            print(f"Epoch {epoch + 1}/{epochs}, Loss: {loss_val:.4f}")

        # 清零梯度，为下一次迭代做准备
        model.zero_grad()

    print("\n训练完成!")
    # e. 测试模型
    test_input = Value(np.array([[6.0]]))
    prediction = model(test_input)
    result = f(6.0)
    print(f"输入 6.0, 模型预测输出: {prediction.data[0, 0]:.4f} (期望值: {result})")


if __name__ == "__main__":
    # 实例化并运行测试
    test_fit()