157 lines
4.7 KiB
Python
157 lines
4.7 KiB
Python
import math
|
||
import random
|
||
import time
|
||
from dataclasses import dataclass
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
TEXT = (
|
||
"起初只有一个简单的故事。"
|
||
"它不长，但足够教会模型去猜下一个字符。"
|
||
"我们会在这段文字上训练一个很小的语言模型，并观察它学习。"
|
||
"实践胜过空谈，所以先从小处开始，不断迭代。"
|
||
"很快我们会用更大的数据替换它。"
|
||
"现在先保持小而清晰。"
|
||
)
|
||
|
||
@dataclass
|
||
class Config:
|
||
# 序列长度（模型一次看到的字符数）
|
||
seq_len: int = 64
|
||
# 批大小
|
||
batch_size: int = 32
|
||
# 字符嵌入维度
|
||
embed_dim: int = 64
|
||
# GRU 隐藏层维度
|
||
hidden_dim: int = 128
|
||
# GRU 层数
|
||
num_layers: int = 1
|
||
# 学习率
|
||
lr: float = 2e-3
|
||
# 训练步数
|
||
steps: int = 800
|
||
# 采样频率
|
||
sample_every: int = 200
|
||
# 采样长度
|
||
sample_len: int = 200
|
||
# 采样温度（越高越随机）
|
||
temperature: float = 0.9
|
||
|
||
|
||
class CharRNN(nn.Module):
|
||
def __init__(self, vocab_size, embed_dim, hidden_dim, num_layers):
|
||
super().__init__()
|
||
# 字符索引 -> 向量
|
||
self.embed = nn.Embedding(vocab_size, embed_dim)
|
||
# GRU 用于处理序列上下文
|
||
self.rnn = nn.GRU(
|
||
embed_dim,
|
||
hidden_dim,
|
||
num_layers=num_layers,
|
||
batch_first=True,
|
||
)
|
||
# 隐藏状态 -> 下一个字符的 logits
|
||
self.fc = nn.Linear(hidden_dim, vocab_size)
|
||
|
||
def forward(self, x, h=None):
|
||
# x: (batch, seq_len)
|
||
x = self.embed(x)
|
||
# out: (batch, seq_len, hidden_dim)
|
||
out, h = self.rnn(x, h)
|
||
# logits: (batch, seq_len, vocab_size)
|
||
logits = self.fc(out)
|
||
return logits, h
|
||
|
||
|
||
def build_vocab(text):
|
||
# 构建字符表与双向映射
|
||
chars = sorted(set(text))
|
||
stoi = {ch: i for i, ch in enumerate(chars)}
|
||
itos = {i: ch for ch, i in stoi.items()}
|
||
return chars, stoi, itos
|
||
|
||
|
||
def encode(text, stoi):
|
||
# 文本 -> 索引序列
|
||
return torch.tensor([stoi[ch] for ch in text], dtype=torch.long)
|
||
|
||
|
||
def get_batch(data, cfg):
|
||
# 随机抽取若干段序列，生成输入 x 和目标 y（y 是 x 向后移一位）
|
||
max_start = len(data) - cfg.seq_len - 1
|
||
starts = torch.randint(0, max_start, (cfg.batch_size,))
|
||
x = torch.stack([data[s : s + cfg.seq_len] for s in starts])
|
||
y = torch.stack([data[s + 1 : s + cfg.seq_len + 1] for s in starts])
|
||
return x, y
|
||
|
||
|
||
def sample(model, stoi, itos, device, cfg, start="起初"):
|
||
# 从给定起始串开始，按概率逐字符采样
|
||
model.eval()
|
||
# 如果起始字符不在词表中，使用文本开头的字符兜底
|
||
safe_start = "".join(ch for ch in start if ch in stoi)
|
||
if not safe_start:
|
||
safe_start = TEXT[:1]
|
||
idx = torch.tensor([stoi[ch] for ch in safe_start], dtype=torch.long, device=device)
|
||
idx = idx.unsqueeze(0)
|
||
h = None
|
||
with torch.no_grad():
|
||
for _ in range(cfg.sample_len):
|
||
logits, h = model(idx, h)
|
||
# 只取最后一个位置的预测
|
||
logits = logits[:, -1, :] / cfg.temperature
|
||
probs = F.softmax(logits, dim=-1)
|
||
next_id = torch.multinomial(probs, num_samples=1)
|
||
idx = torch.cat([idx, next_id], dim=1)
|
||
out = "".join(itos[i] for i in idx[0].tolist())
|
||
return out
|
||
|
||
|
||
def main():
|
||
cfg = Config()
|
||
# 字符表与编码
|
||
chars, stoi, itos = build_vocab(TEXT)
|
||
data = encode(TEXT, stoi)
|
||
|
||
# Apple Silicon 上优先用 MPS
|
||
device = "mps" if torch.backends.mps.is_available() else "cpu"
|
||
model = CharRNN(len(chars), cfg.embed_dim, cfg.hidden_dim, cfg.num_layers).to(device)
|
||
opt = torch.optim.AdamW(model.parameters(), lr=cfg.lr)
|
||
|
||
start_time = time.time()
|
||
model.train()
|
||
for step in range(1, cfg.steps + 1):
|
||
x, y = get_batch(data, cfg)
|
||
x = x.to(device)
|
||
y = y.to(device)
|
||
|
||
logits, _ = model(x)
|
||
# 交叉熵：预测下一个字符
|
||
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1))
|
||
|
||
opt.zero_grad()
|
||
loss.backward()
|
||
# 梯度裁剪，防止梯度爆炸
|
||
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
||
opt.step()
|
||
|
||
if step % 100 == 0:
|
||
elapsed = time.time() - start_time
|
||
speed = step / max(elapsed, 1e-6)
|
||
print(f"step {step} | loss {loss.item():.4f} | {speed:.1f} steps/s")
|
||
|
||
if step % cfg.sample_every == 0:
|
||
print("-" * 60)
|
||
print(sample(model, stoi, itos, device, cfg))
|
||
print("-" * 60)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 固定随机种子，便于复现
|
||
random.seed(42)
|
||
torch.manual_seed(42)
|
||
main()
|
||
|