训练框架总结 - 从零开始大模型

在前十六章，我们从零开始，每一章都是在上一章的基础上逐步扩展，一步步搭建出了张量、计算图、网络层、损失函数、优化器、调度器、训练器……实现了深度学习向大语言模型跨越的基础构件。

现在，是时候把所有这些模块全部收集起来，组装成一个完整的、可复用的网络模型训练框架了。

这一章是前两部分的代码总结；也是下一部分，大语言模型，的起点。我们把现有的基本组件汇集成一个代码库，之后的部分将不再每次重复这些组件的代码。

import contextlib
import math
from abc import ABC, abstractmethod

import numpy as np

DTYPE = np.float32

张量¶

张量是整个框架的基石，也是自动微分的载体。我们对张量进行了一些重要的功能扩展：

运算符重载：我们之前已经实现了张量加法，现在完整实现了减法、乘法、除法，以及矩阵乘法和矩阵转置。使得张量之间可以像数字一样，直接用 +、-、*、/、@ 运算，同时会自动记录计算图，和整个多层模型融为一体；
广播支持：处理不同形状的张量之间的计算，比如一个 (2, 6) 的张量和一个 (1, 6) 的张量的加法。这在处理复杂模型，或者多维度数据时非常常见；
关闭梯度（no_grad）：允许在推理阶段不需要计算图时，临时关闭梯度计算，节约内存和算力；
统一数据类型（DTYPE = np.float32）：统一明确定义张量数值为 32 位浮点数。

⚠️ 注意：运算符的两边必须都是张量。张量和 Python 数字直接运算（例如 x * 2）会报错，需要先把数字包装成张量，例如 x * Tensor(2.0)。

class Tensor:
    grad_enabled = True

    # 上下文管理器：关闭梯度，运行 with 代码块时，暂停计算图、反向传播和参数更新
    @classmethod
    @contextlib.contextmanager
    def no_grad(cls):
        prev = cls.grad_enabled
        cls.grad_enabled = False
        try:
            yield
        finally:
            cls.grad_enabled = prev

    def __init__(self, data):
        self.data = np.asarray(data, dtype=DTYPE)
        self.grad = np.zeros_like(self.data)
        self.gradient_fn = None
        self.parents = []

    def backward(self):
        topo = []
        visited = set()
        stack = [(self, False)]

        while stack:
            node, expanded = stack.pop()
            if node in visited:
                continue

            if expanded:
                visited.add(node)
                topo.append(node)
            else:
                stack.append((node, True))
                for p in node.parents:
                    if p not in visited:
                        stack.append((p, False))

        self.grad = np.ones_like(self.data)
        for t in reversed(topo):
            if t.gradient_fn is not None:
                t.gradient_fn()

        for t in topo:
            t.gradient_fn = None
            t.parents = []

    @property
    def shape(self):
        return self.data.shape

    @property
    def dtype(self):
        return self.data.dtype

    # 张量加法
    def __add__(self, other):
        p = Tensor(self.data + other.data)

        def gradient_fn():
            self.grad += self._unbroadcast(p.grad, self.shape)
            other.grad += self._unbroadcast(p.grad, other.shape)

        return p.attach(gradient_fn, parents=[self, other])

    # 张量减法
    def __sub__(self, other):
        p = Tensor(self.data - other.data)

        def gradient_fn():
            self.grad += self._unbroadcast(p.grad, self.shape)
            other.grad += self._unbroadcast(-p.grad, other.shape)

        return p.attach(gradient_fn, parents=[self, other])

    # 张量乘法
    def __mul__(self, other):
        p = Tensor(self.data * other.data)

        def gradient_fn():
            self.grad += self._unbroadcast(p.grad * other.data, self.shape)
            other.grad += self._unbroadcast(p.grad * self.data, other.shape)

        return p.attach(gradient_fn, parents=[self, other])

    # 张量除法
    def __truediv__(self, other):
        p = Tensor(self.data / other.data)

        def gradient_fn():
            self.grad += self._unbroadcast(p.grad / other.data, self.shape)
            other.grad += self._unbroadcast(-p.grad * self.data / (other.data ** 2), other.shape)

        return p.attach(gradient_fn, parents=[self, other])

    # 张量矩阵乘法（@）
    def __matmul__(self, other):
        p = Tensor(np.matmul(self.data, other.data))

        def gradient_fn():
            self.grad += self._unbroadcast(np.matmul(p.grad, other.data.swapaxes(-1, -2)), self.shape)
            other.grad += self._unbroadcast(np.matmul(self.data.swapaxes(-1, -2), p.grad), other.shape)

        return p.attach(gradient_fn, parents=[self, other])

    # 张量矩阵转置
    def transpose(self, axes=None):
        p = Tensor(np.transpose(self.data, axes))

        def gradient_fn():
            if axes is None:
                self.grad += np.transpose(p.grad)
            else:
                idx = np.argsort(axes)
                self.grad += np.transpose(p.grad, idx)

        return p.attach(gradient_fn, parents=[self])

    # 张量矩阵转置
    @property
    def T(self):
        return self.transpose()

    # 张量变换格式
    def reshape(self, shape):
        p = Tensor(np.reshape(self.data, shape))

        def gradient_fn():
            self.grad += np.reshape(p.grad, self.shape)

        return p.attach(gradient_fn, parents=[self])

    # 设置张量梯度（闭包）函数和父节点列表
    def attach(self, gradient_fn, parents):
        if Tensor.grad_enabled:
            self.gradient_fn = gradient_fn
            self.parents = parents
        return self

    def __repr__(self):
        return f"Tensor(shape={self.shape}, dtype={self.dtype})"

    # 广播机制
    @staticmethod
    def _unbroadcast(grad, shape):
        if grad.ndim > len(shape):
            grad = grad.sum(axis=tuple(range(grad.ndim - len(shape))))

        for axis, dim in enumerate(shape):
            if dim == 1 and grad.shape[axis] != 1:
                grad = grad.sum(axis=axis, keepdims=True)
        return grad.reshape(shape)

模型¶

神经网络模型由多个网络层组成。所有层都支持由抽象层定义的基本功能。

抽象层¶

抽象层定义了最重要的函数接口：

forward（推理函数）：实现前向传播；
parameters（参数列表）：提供所有需要训练的参数；
train（训练函数）：进入训练模式；
eval（测试函数）：进入评估、推理模式。

class Layer(ABC):

    def __init__(self):
        self.training = True

    def __call__(self, *args):
        return self.forward(*args)

    # 训练模式
    def train(self):
        self.training = True

    # 推理模式
    def eval(self):
        self.training = False

    @abstractmethod
    def forward(self, *args):
        pass

    @property
    def parameters(self):
        return []

    def __repr__(self):
        return f"{type(self).__name__}[]"

线性层¶

加权求和，是最基本的网络层。

class Linear(Layer):

    def __init__(self, in_size, out_size):
        super().__init__()
        self.weight = Tensor(np.random.randn(out_size, in_size).astype(DTYPE) * np.sqrt(2 / in_size))
        self.bias = Tensor(np.zeros(out_size).astype(DTYPE))

    def forward(self, x: Tensor):
        p = Tensor(x.data @ self.weight.data.T + self.bias.data)

        def gradient_fn():
            grad = p.grad.reshape(-1, p.grad.shape[-1])
            self.weight.grad += grad.T @ x.data.reshape(-1, x.shape[-1])
            self.bias.grad += np.sum(grad, axis=0)
            x.grad += p.grad @ self.weight.data

        return p.attach(gradient_fn, {x})

    @property
    def parameters(self):
        return [self.weight, self.bias]

组合抽象层¶

定义了可以用作子层容器的抽象层，是复杂网络层的基础。

class Composite(Layer, ABC):

    def __init__(self, layers):
        super().__init__()
        self.layers = list(layers)

    def train(self):
        super().train()
        for l in self.layers:
            l.train()

    def eval(self):
        super().eval()
        for l in self.layers:
            l.eval()

    @property
    def parameters(self):
        return [p for l in self.layers for p in l.parameters]

顺序层¶

最基本的网络层容器，所有子层顺序执行。

class Sequential(Composite):

    def forward(self, x: Tensor):
        for l in self.layers:
            x = l(x)
        return x

嵌入层¶

把字符或者词的索引编码通过查表转换成稠密嵌入向量，避免了使用独热编码的稀疏向量。

class Embedding(Layer):

    def __init__(self, vocab_size, embedding_size, std=0.02):
        super().__init__()
        self.weight = Tensor(np.random.randn(vocab_size, embedding_size).astype(DTYPE) * std)

    def forward(self, x: Tensor):
        p = Tensor(self.weight.data[x.data.astype(np.int64)])

        def gradient_fn():
            np.add.at(self.weight.grad, x.data.astype(np.int64), p.grad)

        return p.attach(gradient_fn, parents=[self.weight])

    @property
    def parameters(self):
        return [self.weight]

平均池化层¶

把一段序列的向量组采用平均值合并成一个向量。

class MeanPool(Layer):

    def forward(self, x: Tensor):
        p = Tensor(np.mean(x.data, axis=1))

        def gradient_fn():
            x.grad += p.grad[:, None, :] / x.data.shape[1]

        return p.attach(gradient_fn, parents=[x])

激活函数¶

激活函数为网络模型引入了非线性，让模型可以学习到复杂关系。

线性整流函数（ReLU）¶

隐藏层默认首选。计算快，不会造成梯度消失。

class ReLU(Layer):

    def forward(self, x: Tensor):
        a = Tensor(np.maximum(0, x.data))

        def gradient_fn():
            x.grad += a.grad * (a.data > 0)

        return a.attach(gradient_fn, parents=[x])

双曲正切函数（Tanh）¶

输出 0 中心化，限制在 (-1, 1) 的范围内，存在梯度饱和问题。

class Tanh(Layer):

    def forward(self, x: Tensor):
        a = Tensor(np.tanh(x.data))

        def gradient_fn():
            x.grad += a.grad * (1 - a.data ** 2)

        return a.attach(gradient_fn, parents=[x])

S 形曲线函数（Sigmoid）¶

用于二元分类问题的输出层。

class Sigmoid(Layer):

    def __init__(self, clip_range=(-100, 100)):
        super().__init__()
        self.clip_range = clip_range

    def forward(self, x: Tensor):
        z = np.clip(x.data, self.clip_range[0], self.clip_range[1])
        a = Tensor(1 / (1 + np.exp(-z)))

        def gradient_fn():
            x.grad += a.grad * a.data * (1 - a.data)

        return a.attach(gradient_fn, parents=[x])

归一化指数函数（Softmax）¶

用于多元分类问题的输出层。

class Softmax(Layer):

    def __init__(self, axis=-1):
        super().__init__()
        self.axis = axis

    def forward(self, x: Tensor):
        exp = np.exp(x.data - np.max(x.data, axis=self.axis, keepdims=True))
        a = Tensor(exp / np.sum(exp, axis=self.axis, keepdims=True))

        def gradient_fn():
            grad = np.sum(a.data * a.grad, axis=self.axis, keepdims=True)
            x.grad += a.data * (a.grad - grad)

        return a.attach(gradient_fn, parents=[x])

损失函数¶

损失函数衡量预测值和标签的误差，是训练的“指挥棒”，反向传播的起始点。

抽象类¶

损失函数抽象类定义了唯一的函数接口：

loss（损失函数）：计算损失值。

class Loss(ABC):

    def __call__(self, *args, **kwargs):
        return self.loss(*args, **kwargs)

    @abstractmethod
    def loss(self, *args, **kwargs):
        pass

均方误差损失函数¶

用于回归问题。

class MSELoss(Loss):

    def loss(self, p: Tensor, y: Tensor):
        mse = Tensor(np.mean(np.square(y.data - p.data)))

        def gradient_fn():
            p.grad += mse.grad * (-2 * (y.data - p.data) / y.data.size)

        return mse.attach(gradient_fn, parents=[p])

交叉熵损失函数¶

用于多元分类问题。内部已经包含 Softmax 计算。

class CELoss(Loss):

    def loss(self, p: Tensor, y: Tensor):
        exp = np.exp(p.data - np.max(p.data, axis=-1, keepdims=True))
        softmax = exp / np.sum(exp, axis=-1, keepdims=True)

        if y.data.ndim == softmax.ndim:
            labels = y.data
        else:
            labels = np.zeros_like(softmax)
            rows = np.arange(y.data.size)
            labels[rows, y.data.astype(np.int64)] = 1

        log = np.log(np.clip(softmax, 1e-10, 1.0))
        ce = Tensor(-np.mean(np.sum(labels * log, axis=-1)))

        def gradient_fn():
            grad = (softmax - labels) / labels.shape[0]
            p.grad += ce.grad * grad

        return ce.attach(gradient_fn, parents=[p])

二元交叉熵损失函数¶

用于二元分类问题。

class BCELoss(Loss):

    def loss(self, p: Tensor, y: Tensor):
        clipped = np.clip(p.data, 1e-7, 1 - 1e-7)
        bce = Tensor(-np.mean(y.data * np.log(clipped) + (1 - y.data) * np.log(1 - clipped)))

        def gradient_fn():
            grad = (clipped - y.data) / (clipped * (1 - clipped)) / y.data.size
            p.grad += bce.grad * grad

        return bce.attach(gradient_fn, parents=[p])

优化器¶

优化器负责根据梯度和学习率更新参数。

抽象类¶

优化器抽象类定义了两个基本函数接口：

step（步进函数）：更新参数；
zero_grad（清零函数）：每次迭代前清空参数梯度。

class Optimizer(ABC):

    def __init__(self, parameters, lr):
        self.parameters = list(parameters)
        self.lr = lr

    @abstractmethod
    def step(self):
        pass

    def zero_grad(self):
        for p in self.parameters:
            p.grad = np.zeros_like(p.data)

随机梯度下降优化器¶

最简单、快速，学习率固定。

class SGDOptimizer(Optimizer):

    def step(self):
        for p in self.parameters:
            p.data -= p.grad * self.lr

自适应矩估计优化器¶

给每个参数单独调整有效步幅，附带梯度裁剪。

class AdamOptimizer(Optimizer):

    def __init__(self, parameters, lr=0.01, betas=(0.9, 0.999), eps=1e-8):
        super().__init__(parameters, lr)
        self.b1, self.b2 = betas
        self.eps = eps
        self.m = [None] * len(parameters)
        self.v = [None] * len(parameters)
        self.t = 0

    def step(self):
        self.t += 1
        for i, p in enumerate(self.parameters):
            if p is not None:
                if self.m[i] is None:
                    self.m[i] = np.zeros_like(p.data)
                    self.v[i] = np.zeros_like(p.data)

                self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * p.grad
                self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * (p.grad ** 2)
                m = self.m[i] / (1 - self.b1 ** self.t)
                v = self.v[i] / (1 - self.b2 ** self.t)
                self._apply_weight_decay(p)
                p.data -= self.lr * m / (np.sqrt(v) + self.eps)

    def _apply_weight_decay(self, p):
        pass

    def clip_grad_norm(self, max_norm=1.0):
        sq = 0.0
        for p in self.parameters:
            sq += float(np.sum(p.grad ** 2))

        if np.sqrt(sq) > max_norm > 0:
            scale = max_norm / (np.sqrt(sq) + 1e-6)
            for p in self.parameters:
                p.grad *= scale

权重衰减自适应矩估计优化器¶

附带权重衰减的 Adam。

class AdamWOptimizer(AdamOptimizer):

    def __init__(self, parameters, lr=0.01, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01):
        super().__init__(parameters, lr, betas, eps)
        self.weight_decay = weight_decay

    def _apply_weight_decay(self, p):
        if p.data.ndim >= 2:
            p.data -= p.data * self.weight_decay * self.lr

学习率调度器¶

学习率调度器可以在训练过程中动态调整学习率（自动调参）。

抽象类¶

学习率调度器抽象类定义唯一的函数接口：

step（步进函数）：根据当前步数调整学习率。

class LRScheduler(ABC):

    @abstractmethod
    def step(self, current_step):
        pass

预热余弦退火调度器¶

让学习率从 0 爬升（预热），再沿余弦曲线平滑下降（余弦退火）。是训练大模型的标准做法。

class WarmupCosineScheduler(LRScheduler):

    def __init__(self, lr, total_steps, warmup_steps, min_lr=0.0):
        self.lr = lr
        self.total_steps = max(total_steps, 1)
        self.warmup_steps = max(min(warmup_steps, self.total_steps), 0)
        self.min_lr = min_lr

    def step(self, current_step):
        if self.warmup_steps > 0 and current_step < self.warmup_steps:
            return self.lr * (current_step + 1) / self.warmup_steps

        if current_step >= self.total_steps:
            return self.min_lr

        progress = (current_step - self.warmup_steps) / max(self.total_steps - self.warmup_steps, 1)
        cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
        return self.min_lr + (self.lr - self.min_lr) * cosine

数据集¶

数据集负责加载数据，划分训练集、测试集，以及按照批次遍历所有样本。

抽象类¶

数据集抽象类定义了基本的函数接口：

load（加载函数）：加载数据；
train（训练函数）：提供训练集数据；
eval（测试函数）：提供测试集数据；
all（全部样本）：读取全部样本；
__len__（样本数量）：返回样本数量；
__getitem__（单个样本）：根据索引读取样本。

class Dataset(ABC):

    def __init__(self, batch_size=1):
        self.batch_size = batch_size
        self.load()
        self.train()

    @abstractmethod
    def load(self):
        pass

    def train(self):
        self.data = self.train_data

    def eval(self):
        self.data = self.test_data

    def all(self):
        x, y = self.data
        return Tensor(x), Tensor(y)

    def __len__(self):
        x, *_ = self.data
        return len(x) // self.batch_size

    def __getitem__(self, index):
        s = slice(index * self.batch_size, (index + 1) * self.batch_size)
        x, y = self.data
        return Tensor(x[s]), Tensor(y[s])

训练器¶

训练器把模型、损失函数、优化器打包在一起，实现模型训练和测试逻辑。

抽象类¶

训练器抽象类定义两个基本的函数接口：

train（训练函数）：用数据集训练集进行多轮次训练，可以接入学习率调度器；
test（测试函数）：在数据集测试集上评估模型。

class Trainer(ABC):

    def __init__(self, model, loss_fn, optimizer):
        self.model = model
        self.loss_fn = loss_fn
        self.optimizer = optimizer

    @abstractmethod
    def train(self, dataset, epochs, scheduler=None):
        pass

    @abstractmethod
    def test(self, dataset):
        pass

至此，我们的核心框架基本成型。

这套框架虽然只用 NumPy 写成，但是具备了现代深度学习框架的核心功能。下一部分，我们将在此基础上，搭建真正的 GPT 模型。

课后练习¶

张量目前不支持和 Python 数字直接运算。请修改 __add__、__sub__、__mul__、__truediv__，让 x * 2 和 x + 1.0 也能正常运行，并且反向传播的梯度仍然正确。提示：在运算开始时，先把数字包装成张量。