学习率调度器 - 从零开始大模型

上一章，Adam 帮我们给每个参数单独地调整有效步幅。与步幅相关的还有一个超参数：学习率，我们仍旧是手动调参。

目前的学习率，在开始训练后是固定的。训练开始时是多少，一直到结束还是多少。但是这并不合理：

训练初期：参数还是随机初始化的，梯度方向不可靠。如果一上来就用大学习率猛冲，容易直接跑偏；
训练后期：模型已经接近最优解，需要小心微调，而不是继续大步走，否则会在最优点附近来回震荡，甚至跳脱。

这一章，我们将引入自动调参：用学习率调度器（Learning Rate Scheduler）在训练过程中动态地改变学习率。我们将采用大模型训练最常用的方案：预热（Warmup）和余弦退火（Cosine Annealing）。

预热：训练初期，学习率从 0 开始，逐步爬升到最大值；
余弦退火：预热结束后，学习率沿着余弦曲线，从最大值平滑地下降到最小值。

import math
from abc import ABC, abstractmethod

import numpy as np

np.random.seed(42)

张量¶

class Tensor:

    def __init__(self, data):
        self.data = np.array(data)
        self.grad = np.zeros_like(self.data)
        self.gradient_fn = None
        self.parents = set()

    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 = lambda: None
            t.parents = set()

    def __add__(self, other):
        p = Tensor(self.data + other.data)

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

        p.gradient_fn = gradient_fn
        p.parents = [self, other]
        return p

    def __str__(self):
        return f'Tensor({self.data})'

数据集¶

抽象类¶

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])

字符编码数据集¶

class CharDataset(Dataset):

    def __init__(self, filename, batch_size=1, context_size=64, stride=None, split=0.9):
        self.filename = filename
        self.context_size = context_size
        self.stride = stride if stride is not None else context_size // 2
        self.split = split
        super().__init__(batch_size)

    def load(self):
        with open(self.filename, encoding="utf-8") as f:
            text = f.read()

        self.vocab = sorted(set(text))
        self.vocab_size = len(self.vocab)
        self.stoi = {ch: i for i, ch in enumerate(self.vocab)}
        self.itos = {i: ch for i, ch in enumerate(self.vocab)}
        self.tokens = self.encode(text)

        split = int(len(self.tokens) * self.split)
        self.train_data = self._pack(self.tokens[:split])
        self.test_data = self._pack(self.tokens[split:])

    def _pack(self, tokens):
        x, y = [], []
        for i in range(0, len(tokens) - self.context_size - 1, self.stride):
            x.append(tokens[i: i + self.context_size])
            y.append(tokens[i + 1: i + self.context_size + 1])
        return x, y

    def encode(self, symbols):
        return [self.stoi[s] for s in symbols]

    def decode(self, tokens):
        return "".join(self.itos[t] for t in tokens)

模型¶

抽象层¶

class Layer(ABC):

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

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

    @property
    def parameters(self):
        return []

线性层¶

class Linear(Layer):

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

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

        def gradient_fn():
            self.weight.grad += p.grad.T @ x.data
            self.bias.grad += np.sum(p.grad, axis=0)
            x.grad += p.grad @ self.weight.data

        p.gradient_fn = gradient_fn
        p.parents = [x]
        return p

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

组合抽象层¶

class Composite(Layer, ABC):

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

    @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) * std)

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

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

        p.gradient_fn = gradient_fn
        p.parents = [self.weight]
        return p

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

激活函数¶

双曲正切函数（Tanh）¶

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)

        a.gradient_fn = gradient_fn
        a.parents = [x]
        return a

损失函数¶

抽象类¶

class Loss(ABC):

    def __call__(self, p: Tensor, y: Tensor):
        return self.loss(p, y)

    @abstractmethod
    def loss(self, p: Tensor, y: Tensor):
        pass

交叉熵损失函数¶

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] = 1

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

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

        ce.gradient_fn = gradient_fn
        ce.parents = [p]
        return ce

优化器¶

抽象类¶

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 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

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

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

学习率调度器（预热余弦退火）¶

创建一个学习率调度器，我们需要提供 4 个初始化参数：

lr（学习率）：这就是我们一直使用的学习率，也是预热后达到的最大学习率；
total_steps（总步数）：整个训练过程的全部步数；
warmup_steps（预热步数）：预热阶段的步数；
min_lr（最小学习率）：训练后期学习率降到的最小值。

预热¶

训练的最开始 warmup_steps 步，学习率从 0 线性爬升到最大值 lr：

lr_t = lr \cdot \frac{t}{warmup\_steps}

(1)

这里：

$lr_t$ ：当前学习率；
$t$ ：当前步数（从 1 开始计数，对应代码里的 current_step + 1）。

余弦退火¶

预热结束后，学习率达到最大值，然后沿着余弦曲线，从最大值平滑地下降到最小值 $min\_lr$ ：

lr_t = min\_lr + (lr - min\_lr) \cdot \frac{1 + \cos(\pi \cdot \frac{t - warmup\_steps}{total\_steps - warmup\_steps})}{2}

(2)

这里：

$lr_t$ ：当前学习率；
$t$ ：当前步数（对应代码里的 current_step）。

选用余弦曲线来退火的好处是：步幅先慢、后快、再慢。所以整个训练周期，学习率就像开车：慢慢起步（预热），之后中段加速巡航，快到终点时平稳减速（余弦退火）。这是 GPT 这类大模型的标准训练节奏。

学习率调度器只提供一个函数：

step（步进函数）：根据当前步数计算出新的学习率。

class WarmupCosineScheduler:

    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

循环神经网络¶

循环单元¶

class Cell(Composite):

    def __init__(self, in_size, out_size):
        self.input = Linear(in_size, out_size)
        self.hidden = Linear(out_size, out_size)
        self.tanh = Tanh()

        super().__init__([self.input,
                          self.hidden,
                          self.tanh])

    def forward(self, x: Tensor, h: Tensor):
        p = self.input(x)
        h = self.hidden(h)
        return self.tanh(p + h)

循环神经网络模型¶

class RNN(Composite):

    def __init__(self, vocab_size, embedding_size):
        self.embedding_size = embedding_size

        self.embedding = Embedding(vocab_size, embedding_size)
        self.cell = Cell(embedding_size, embedding_size * 2)
        self.output = Linear(embedding_size * 2, vocab_size)

        super().__init__([self.embedding,
                          self.cell,
                          self.output])

    def step(self, token: Tensor, h: Tensor):
        p = self.embedding(token)
        h = self.cell(p, h)
        return self.output(h), h

    def forward(self, x: Tensor, h: Tensor = None):
        batch_size, sequence_len = x.data.shape
        if h is None:
            h = Tensor(np.zeros((batch_size, self.embedding_size * 2)))

        ps = []
        for s in range(sequence_len):
            p, h = self.step(Tensor(x.data[:, s]), h)
            ps.append(p)

        return ps, h

训练器¶

抽象类¶

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

循环神经网络训练器¶

训练器中每一步都先调用学习率调度器的步进函数 step，获得最新的学习率。

class RNNTrainer(Trainer):

    def train(self, dataset, epochs, scheduler=None):
        dataset.train()

        steps = 0
        for epoch in range(epochs):
            for i in range(len(dataset)):
                # 更新学习率
                if scheduler is not None:
                    self.optimizer.lr = scheduler.step(steps)

                feature, label = dataset[i]

                self.optimizer.zero_grad()
                prediction, _ = self.model(feature)
                loss = Tensor(0.0)
                for t in range(len(prediction)):
                    loss += self.loss_fn(prediction[t], Tensor(label.data[:, t]))
                loss.backward()
                self.optimizer.clip_grad_norm()
                self.optimizer.step()
                steps += 1

    def test(self, dataset):
        dataset.eval()

        feature, label = dataset.all()
        prediction, _ = self.model(feature)
        loss = Tensor(0.0)
        for t in range(len(prediction)):
            loss += self.loss_fn(prediction[t], Tensor(label.data[:, t]))
        return prediction, loss

    def generate(self, dataset, prompt, steps=256):
        tokens = dataset.encode(prompt)

        h = Tensor(np.zeros((1, self.model.embedding_size * 2)))
        for token in tokens:
            p, h = self.model.step(Tensor([token]), h)

        for _ in range(steps):
            exp = np.exp(p.data[0] - np.max(p.data[0]))
            probs = exp / np.sum(exp)
            token = np.random.choice(len(probs), p=probs)
            tokens.append(token)
            p, h = self.model.step(Tensor([token]), h)

        return dataset.decode(tokens)

超参数¶

学习率¶

使用学习率调度器后，LEARNING_RATE 实际上用作最大学习率。学习率调度器降低平均学习率，所以我们可以适当调大一些，而最小学习率设为最大学习率的 1/10。

LEARNING_RATE = 0.005

批大小¶

BATCH_SIZE = 4

轮数¶

EPOCHS = 2

上下文尺寸¶

CONTEXT_SIZE = 32

嵌入尺寸¶

EMBEDDING_SIZE = 16

建模¶

DATA_FILE = "../../../../data/tinyshakespeare.txt"

dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = RNN(dataset.vocab_size, EMBEDDING_SIZE)
loss_fn = CELoss()
optimizer = AdamWOptimizer(model.parameters, lr=LEARNING_RATE)
trainer = RNNTrainer(model, loss_fn, optimizer)

训练¶

开始模型训练前，我们要先创建一个学习率调度器的实例，并把它传给训练函数。这里几个初始化参数的含义是：

LEARNING_RATE：最大学习率；
EPOCHS * len(dataset)：全部训练步数；
100：预热步数；
LEARNING_RATE / 10：最小学习率。

# 学习率调度器
scheduler = WarmupCosineScheduler(LEARNING_RATE, EPOCHS * len(dataset), 100, LEARNING_RATE / 10)
trainer.train(dataset, EPOCHS, scheduler)

评估¶

prediction, loss = trainer.test(dataset)
print(f'prediction:\t{len(prediction)} steps, each {prediction[0].data.shape}\ntotal loss:\t{loss}')

prediction:	32 steps, each (6970, 65)
total loss:	Tensor(64.7635956614507)

推理¶

print(trainer.generate(dataset, prompt="ROMEO:"))

ROMEO:
Hat that, he you!
You spelr gexssigean I'red
ho,
That his whirovellm to waind:
Ye Gosis
The rest patiet,
Pervunded mantlacem! But do stay,
Hath muthood;
I wore not hanted gaing, you knes is not thought be but what you a you geeden Hartion as bearter a Pro

这一章，我们实现了学习率的自动调参：预热让训练平稳起步，余弦退火让训练收尾精细收敛。至此，模型训练的三大实战神器：Adam 优化器、梯度裁剪、学习率调度，都配齐了，这也是训练大模型的标准配置。

回顾第二部分，我们从数字化文字信息开始，实现了 RNN，创建了一个可以读懂文字的模型。

但是必须指出的是：RNN 存在着记忆有限、无法并行的先天缺陷，无法承担构建真正大语言模型的重担。下一部分，我们将换一种全新的思路：注意力机制（Attention），并一步步搭建出真正的 GPT 模型。

课后练习¶

试着调整学习率调度器的初始化参数，比如把预热步数改成 0（不要预热），或者很大；或者把最小学习率改成 0，让学习率最终降到 0。观察训练速度和效果变好还是变差？