上一章，我们给 GPT 模型添加了随机失活的能力，预训练的三大工程优化技术已经完备。
现在，我们来考虑模型预训练的最后一个重要问题：怎样保存训练结果？
训练一个 GPT 模型动则几个小时，甚至数天。显然，不可能每次使用时都重新训练。我们需要把训练好的模型保存下来，下次可以直接加载使用。
保存什么¶
模型的全部知识都在可训练参数里面：每一层的权重和偏置。我们只需要保存这些数值，不需要保存代码、计算图、或者梯度等。
保存格式¶
NumPy 提供了 np.savez 函数，可以把多个数组保存到单一文件（.npz 格式的归档文件）中。
GPT 模型的所有可训练参数全部存储在模型的参数列表 model.parameters 里，我们只需要依次将其中每个参数保存到文件中。
何时保存¶
训练的过程中，我们可以在每个轮次结束后，评估模型在测试集上的表现。如果当前轮次在测试集上的损失值是创了新低，就把模型保存下来。这样，我们可以确保留下了最佳训练结果。
怎么加载¶
下次使用模型时，我们可以从 .npz 文件中读取参数，依次赋值给模型参数 model.parameters。
import os
import numpy as np
from src.core import *np.random.seed(42)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 Tril(Layer):
def __init__(self, value=-1e9):
super().__init__()
self.value = value
def forward(self, x: Tensor):
keep = np.tril(np.ones(x.shape[-2:]))
p = Tensor(np.where(keep, x.data, self.value))
def gradient_fn():
x.grad += p.grad * keep
return p.attach(gradient_fn, {x})归一化层¶
class LayerNorm(Layer):
def __init__(self, normalized_size, eps=1e-5):
super().__init__()
self.eps = eps
self.weight = Tensor(np.ones(normalized_size, dtype=DTYPE))
self.bias = Tensor(np.zeros(normalized_size, dtype=DTYPE))
def forward(self, x: Tensor):
mean = np.mean(x.data, axis=-1, keepdims=True)
var = np.var(x.data, axis=-1, keepdims=True, ddof=0)
norm = (x.data - mean) / np.sqrt(var + self.eps)
p = Tensor(self.weight.data * norm + self.bias.data)
def gradient_fn():
axis = tuple(range(p.grad.ndim - 1)) if p.grad.ndim > 1 else None
self.weight.grad += np.sum(p.grad * norm, axis=axis)
self.bias.grad += np.sum(p.grad, axis=axis)
grad = p.grad * self.weight.data
grad_mean = np.mean(grad, axis=-1, keepdims=True)
norm_mean = np.mean(grad * norm, axis=-1, keepdims=True)
x.grad += (grad - grad_mean - norm * norm_mean) / np.sqrt(var + self.eps)
return p.attach(gradient_fn, {x})
@property
def parameters(self):
return [self.weight, self.bias]随机失活层¶
class Dropout(Layer):
def __init__(self, prob=0.1):
super().__init__()
self.prob = prob
def forward(self, x: Tensor):
if not self.training or self.prob == 0:
return x
keep_prob = 1.0 - self.prob
mask = (np.random.rand(*x.shape) < keep_prob).astype(DTYPE) / keep_prob
p = Tensor(x.data * mask)
def gradient_fn():
x.grad += p.grad * mask
return p.attach(gradient_fn, {x})class GELU(Layer):
def __init__(self):
super().__init__()
self.c = np.sqrt(2.0 / np.pi)
def forward(self, x: Tensor):
tanh = np.tanh(self.c * (x.data + 0.044715 * x.data ** 3))
a = Tensor(0.5 * x.data * (1.0 + tanh))
def gradient_fn():
grad = 0.5 * (1.0 + tanh) + 0.5 * x.data * (1.0 - tanh ** 2) * self.c * (1.0 + 3.0 * 0.044715 * x.data ** 2)
x.grad += a.grad * grad
return a.attach(gradient_fn, parents=[x])class GPTLoss(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)
flat_softmax = softmax.reshape(-1, softmax.shape[-1])
flat_y = y.data.reshape(-1).astype(np.int64)
n = len(flat_y)
rows = np.arange(n)
log = np.log(np.clip(flat_softmax[rows, flat_y], 1e-10, 1))
ce = Tensor(0 - np.sum(log) / n)
def gradient_fn():
flat_grad = flat_softmax.copy()
flat_grad[rows, flat_y] -= 1
p.grad += ce.grad * flat_grad.reshape(softmax.shape) / n
return ce.attach(gradient_fn, parents=[p])class GPTEmbedding(Composite):
def __init__(self, vocab_size, context_size, embedding_size):
self.embedding = Embedding(vocab_size, embedding_size)
self.positional_embedding = Embedding(context_size, embedding_size)
self.dropout = Dropout()
super().__init__([self.embedding,
self.positional_embedding,
self.dropout])
def forward(self, x: Tensor):
token = self.embedding(x)
position = self.positional_embedding(Tensor(range(x.shape[1])))
return self.dropout(token + position)GPT 注意力层¶
class GPTAttention(Composite):
def __init__(self, embedding_size, heads=1):
self.embedding_size = embedding_size
self.heads = heads
self.normalize = LayerNorm(embedding_size)
self.query = Linear(embedding_size, embedding_size)
self.key = Linear(embedding_size, embedding_size)
self.value = Linear(embedding_size, embedding_size)
self.mask = Tril()
self.softmax = Softmax()
self.output = Linear(embedding_size, embedding_size)
self.dropout = Dropout()
super().__init__([self.normalize,
self.query,
self.key,
self.value,
self.mask,
self.softmax,
self.output,
self.dropout])
def forward(self, x: Tensor):
norm = self.normalize(x)
head_embedding_size = self.embedding_size // self.heads
multi_head_shape = (-1, x.shape[1], self.heads, head_embedding_size)
shape = (-1, x.shape[1], self.embedding_size)
query = self.query(norm).reshape(multi_head_shape).transpose((0, 2, 1, 3))
key = self.key(norm).reshape(multi_head_shape).transpose((0, 2, 3, 1))
value = self.value(norm).reshape(multi_head_shape).transpose((0, 2, 1, 3))
scale = Tensor(np.array(1.0 / np.sqrt(head_embedding_size)))
scores = query @ key * scale
weights = self.softmax(self.mask(scores))
return x + self.dropout(self.output((weights @ value).transpose((0, 2, 1, 3)).reshape(shape)))GPT 前馈网络层¶
class GPTFeedForward(Composite):
def __init__(self, embedding_size):
self.normalize = LayerNorm(embedding_size)
self.input = Linear(embedding_size, embedding_size * 4)
self.gelu = GELU()
self.output = Linear(embedding_size * 4, embedding_size)
self.dropout = Dropout()
super().__init__([self.normalize,
self.input,
self.gelu,
self.output,
self.dropout])
def forward(self, x: Tensor):
norm = self.normalize(x)
h = self.gelu(self.input(norm))
return x + self.dropout(self.output(h))GPT Transformer 模块¶
class GPTTransformer(Composite):
def __init__(self, embedding_size, heads):
self.attention = GPTAttention(embedding_size, heads)
self.feed_forward = GPTFeedForward(embedding_size)
super().__init__([self.attention,
self.feed_forward])
def forward(self, x: Tensor):
x = self.attention(x)
return self.feed_forward(x)GPT 输出层¶
class GPTOutput(Composite):
def __init__(self, embedding_size, vocab_size):
self.normalize = LayerNorm(embedding_size)
self.output = Linear(embedding_size, vocab_size)
super().__init__([self.normalize,
self.output])
def forward(self, x: Tensor):
norm = self.normalize(x)
return self.output(norm)GPT 模型¶
class GPT(Composite):
def __init__(self, vocab_size, context_size, embedding_size, heads, blocks):
self.embedding = GPTEmbedding(vocab_size, context_size, embedding_size)
self.transformers = [GPTTransformer(embedding_size, heads) for _ in range(blocks)]
self.output = GPTOutput(embedding_size, vocab_size)
super().__init__([self.embedding] + self.transformers + [self.output])
def forward(self, x: Tensor):
x = self.embedding(x)
for layer in self.transformers:
x = layer(x)
return self.output(x)GPT 训练器¶
我们在 GPT 训练器中添加保存和加载模型参数的功能：
save（保存函数）：将模型所有参数按照顺序编号，保存到.npz文件；load（加载函数）：从.npz文件加载数据，按照顺序赋值给模型的参数；train（训练函数）：重排：每个轮次开始前，我们添加了数据重排功能。打乱样本顺序，增加训练的随机性；
日志：跟踪训练中损失值和学习率的变化；
保存最佳模型参数：每个轮次完成后，用测试集评估训练结果。如果效果更佳，则替代已保存参数。
class GPTTrainer(Trainer):
def train(self, dataset, epochs, scheduler=None, filename=None):
dataset.train()
self.model.train()
steps = 0
# 最佳损失值
best_losses = float('inf')
for epoch in range(epochs):
order = list(range(len(dataset)))
# 重排
np.random.shuffle(order)
total_loss = 0.0
for step, i in enumerate(order):
if scheduler is not None:
self.optimizer.lr = scheduler.step(steps)
feature, label = dataset[i]
prediction = self.model(feature)
loss = self.loss_fn(prediction, label)
self.optimizer.zero_grad()
loss.backward()
total_loss += float(loss.data)
self.optimizer.clip_grad_norm()
self.optimizer.step()
steps += 1
# 训练日志
if (step + 1) % 1000 == 0:
lr = f" lr {self.optimizer.lr:.6f}" if scheduler is not None else ""
print(f"epoch {epoch + 1} step {step + 1}/{len(dataset)} loss {(total_loss / 1000):.4f}{lr}")
total_loss = 0.0
_, total_loss = self.test(dataset)
# 测试日志
print(f"epoch {epoch + 1} done, test loss {total_loss:.4f}")
# 保存最佳模型参数
if total_loss < best_losses and filename is not None:
best_losses = total_loss
self.save(filename)
print(f"epoch {epoch + 1} saved model to {filename}")
print("GPT-training completed")
def test(self, dataset):
dataset.eval()
self.model.eval()
predictions = []
total_loss = 0.0
with Tensor.no_grad():
for i in range(len(dataset)):
feature, label = dataset[i]
prediction = self.model(feature)
loss = self.loss_fn(prediction, label)
predictions.append(prediction)
total_loss += float(loss.data)
loss = total_loss / len(dataset)
dataset.train()
self.model.train()
return predictions, loss
def generate(self, dataset, prompt, steps=256):
self.model.eval()
tokens = dataset.encode(prompt)
with Tensor.no_grad():
for _ in range(steps):
feature = Tensor([tokens[-dataset.context_size:]])
logits = self.model(feature)
last_logits = logits.data[0, -1]
exp = np.exp(last_logits - np.max(last_logits))
probs = exp / np.sum(exp)
token = np.random.choice(len(probs), p=probs)
tokens.append(token)
return dataset.decode(tokens)
# 保存函数
def save(self, filename):
params = {f"param_{i}": p.data for i, p in enumerate(self.model.parameters)}
np.savez(filename, **params)
# 加载函数
def load(self, filename):
if os.path.isfile(filename):
data = np.load(filename, allow_pickle=False)
for i, p in enumerate(self.model.parameters):
p.data = data[f"param_{i}"]
p.grad = np.zeros_like(p.data)LEARNING_RATE = 0.001批大小¶
BATCH_SIZE = 16轮数¶
EPOCHS = 10上下文尺寸¶
CONTEXT_SIZE = 256嵌入尺寸¶
EMBEDDING_SIZE = 384头数¶
HEADS = 6Transformer 模块数¶
BLOCKS = 6建模¶
DATA_FILE = "../../../../data/tinyshakespeare.txt"
dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = GPT(dataset.vocab_size, CONTEXT_SIZE, EMBEDDING_SIZE, HEADS, BLOCKS)
loss_fn = GPTLoss()
optimizer = AdamWOptimizer(model.parameters, lr=LEARNING_RATE)
trainer = GPTTrainer(model, loss_fn, optimizer)训练¶
训练时传入 MODEL_FILE，每个轮次结束后自动保存最佳模型参数。
MODEL_FILE = "../../../../data/tinyshakespeare-gpt.npz"
scheduler = WarmupCosineScheduler(LEARNING_RATE, EPOCHS * len(dataset), 100, LEARNING_RATE / 10)
# 训练时自动保存最佳模型参数
trainer.train(dataset, EPOCHS, scheduler, MODEL_FILE)epoch 1 done, test loss 2.3235
epoch 1 saved model to ../../../../data/tinyshakespeare-gpt.npz
epoch 2 done, test loss 2.0444
epoch 2 saved model to ../../../../data/tinyshakespeare-gpt.npz
epoch 3 done, test loss 1.8580
epoch 3 saved model to ../../../../data/tinyshakespeare-gpt.npz
epoch 4 done, test loss 1.7599
epoch 4 saved model to ../../../../data/tinyshakespeare-gpt.npz
epoch 5 done, test loss 1.6474
epoch 5 saved model to ../../../../data/tinyshakespeare-gpt.npz
epoch 6 done, test loss 1.5968
epoch 6 saved model to ../../../../data/tinyshakespeare-gpt.npz
epoch 7 done, test loss 1.5635
epoch 7 saved model to ../../../../data/tinyshakespeare-gpt.npz
epoch 8 done, test loss 1.5383
epoch 8 saved model to ../../../../data/tinyshakespeare-gpt.npz
epoch 9 done, test loss 1.5339
epoch 9 saved model to ../../../../data/tinyshakespeare-gpt.npz
epoch 10 done, test loss 1.5331
epoch 10 saved model to ../../../../data/tinyshakespeare-gpt.npz
GPT-training completed
评估¶
prediction, loss = trainer.test(dataset)
print(f'prediction:\t{len(prediction)} steps, each {prediction[0].data.shape}\nloss:\t{loss}')prediction: 54 steps, each (16, 256, 65)
loss: 1.5330633212018896
推理¶
训练完成后，最佳模型参数已经保存到 MODEL_FILE 文件。下次使用时，可以直接加载，无需重新训练。
# 加载模型参数
trainer.load(MODEL_FILE)
print(trainer.generate(dataset, prompt="ROMEO:"))ROMEO:
As have done!
LEONTES:
Ho! when be letimest gentleman learness?
PAULINA:
Preventage straight the understand; such nor
Your victory fallse to dare.
DUKE VINCENTIO:
There, by the crown.
GLOUCESTER:
Indeed: thou art uncle and Verona wilt
Of thy sun-she s
本章，我们给 GPT 训练器添加了保存和加载模型参数的功能：训练时自动保存最佳模型的参数，下次可以直接加载使用。
至此，GPT 模型预训练部分就完整了。下一阶段，我们将进入推理（Inference）阶段：怎样使用预训练好的模型生成文本。
课后练习¶
试着把轮数
EPOCHS改成 20，训练更长时间。观察保存的模型参数文件大小和文本生成的质量有什么变化？