File size: 4,425 Bytes
5c6c847 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | """各个实验脚本共用的部分：语料、批次、训练循环、计算量统计、结果存档。"""
import datetime
import glob
import os
import sys
import sysconfig
import time
import numpy as np
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, ".."))
import zerocg as zc # noqa: E402
from models import ZeroX # noqa: E402
from zerocg import functional as F # noqa: E402
ITERS = int(os.environ.get("ITERS", 2000))
T, B = 64, 32
CFG = dict(n_embd=64, n_head=4, n_layer=2)
def load_corpus(path=None):
"""默认语料：Python 标准库源码（纯 ASCII 的 .py 文件，约 150 万字符）。"""
if path:
return open(path, encoding="utf-8").read()
parts, total = [], 0
for p in sorted(glob.glob(os.path.join(sysconfig.get_paths()["stdlib"], "*.py"))):
try:
s = open(p, encoding="utf-8").read()
except (UnicodeDecodeError, OSError):
continue
if s.isascii():
parts.append(s)
total += len(s)
if total > 1_500_000:
break
return "\n".join(parts)
def prepare_data(path=None):
text = load_corpus(path)
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
data = np.array([stoi[c] for c in text], dtype=np.int64)
split = int(0.9 * len(data))
return text, len(chars), data[:split], data[split:]
def batches(data, n, length, seed, batch=B):
r = np.random.default_rng(seed)
for _ in range(n):
ix = r.integers(0, len(data) - length - 1, batch)
yield np.stack([data[i:i + length] for i in ix]), np.stack([data[i + 1:i + length + 1] for i in ix])
def loss_of(logits, y):
return F.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1)).item()
def train(name, model, train_data, val_data, seed, lr=3e-3, iters=None, seq_len=None, batch=None):
"""返回验证 loss（固定在 64 字上评估，便于各实验比较）。
学习率：100 步线性预热 + 余弦退火到 10%；AdamW，权重衰减 0.01，梯度裁剪 1.0。
seq_len / batch 默认 T / B；改训练长度时请同时改 batch，保持每步的字数不变。"""
iters = iters or ITERS
opt = zc.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
t0 = time.time()
for it, (x, y) in enumerate(batches(train_data, iters, seq_len or T, seed, batch or B)):
opt.lr = lr * min(1, (it + 1) / 100) * (0.1 + 0.9 * 0.5 * (1 + np.cos(np.pi * it / iters)))
logits = model(x)
loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1))
opt.zero_grad()
loss.backward()
zc.optim.clip_grad_norm(model.parameters(), 1.0)
opt.step()
if it % 500 == 0:
print(f" [{name} seed={seed} lr={lr:g}] step {it:5d} train {loss.item():.3f} "
f"{time.time() - t0:.0f}s", flush=True)
return evaluate(model, val_data)
def evaluate(model, val_data, n=40):
model.eval()
with zc.no_grad():
val = float(np.mean([loss_of(model(x), y) for x, y in batches(val_data, n, T, 999)]))
model.train()
return val
def forward_flops_per_char(model, V):
"""前向一次、平均每个字多少次浮点运算（只算矩阵乘）。训练一步约是它的 3 倍。"""
x = np.zeros((1, T), dtype=np.int64)
with zc.no_grad(), zc.count_flops() as c:
model(x)
return c.flops / T
def slot_usage(model):
"""每层"有效槽数"：最近一次 forward 的平均读/写地址分布的 exp(熵)，对头取平均。
16 个槽都用得一样多 = 16；全挤在一个槽 = 1。"""
out = []
for b in model.blocks:
read, write = b.mix.usage
eff = lambda p: float(np.mean(np.exp(-(p * np.log(p + 1e-12)).sum(-1)))) # noqa: E731
out.append((eff(read), eff(write)))
return out
def save_result(name, lines):
"""结果写两份：experiments/results/<时间>-<name>.md（存档，不覆盖）和 experiments/<name>.md（最新）。"""
text = "\n".join(lines) + "\n"
stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M")
os.makedirs(os.path.join(HERE, "results"), exist_ok=True)
for path in (os.path.join(HERE, "results", f"{stamp}-{name}.md"), os.path.join(HERE, f"{name}.md")):
open(path, "w").write(text)
print("\n" + text)
def is_zeropoint(model):
return isinstance(model, ZeroX)
|