"""框架自检：用数值差分验证每个算子 / 层的反向传播，并跑两个小模型确认能学会。 运行: python3 tests/test_zerocg.py """ import os import sys import numpy as np sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")) import zerocg as zc # noqa: E402 from zerocg import functional as F # noqa: E402 from zerocg import nn # noqa: E402 zc.set_default_dtype(np.float64) # 数值梯度检查需要 float64 的精度 rng = np.random.default_rng(0) def grad_check(fn, *shapes, tol=1e-5, positive=False): """fn(*tensors) -> Tensor。比较自动求导和中心差分的梯度（float64）。""" arrays = [rng.standard_normal(s) for s in shapes] if positive: arrays = [np.abs(a) + 0.5 for a in arrays] ts = [zc.Tensor(a.copy(), requires_grad=True) for a in arrays] w = np.asarray(rng.standard_normal(fn(*ts).shape)) # 随机投影成标量，检查完整雅可比 (fn(*ts) * zc.Tensor(w)).sum().backward() eps = 1e-6 for t, a in zip(ts, arrays): num = np.zeros_like(a) for i in np.ndindex(a.shape): for sign in (1, -1): b = a.copy() b[i] += sign * eps inputs = [zc.Tensor(b) if x is t else zc.Tensor(x.data) for x in ts] with zc.no_grad(): num[i] += sign * (fn(*inputs).data * w).sum() num[i] /= 2 * eps err = np.abs(num - t.grad).max() / max(1.0, np.abs(num).max()) assert err < tol, f"梯度误差 {err:.2e}" def module_grad_check(module, x): for p in module.parameters(): p.data = p.data.astype(np.float64) params = module.parameters() # 直接对参数逐个做差分 xt = zc.Tensor(x) w = np.asarray(rng.standard_normal(module(xt).shape)) for p in params: p.grad = None (module(xt) * zc.Tensor(w)).sum().backward() eps = 1e-6 for p in params: num = np.zeros_like(p.data) for i in np.ndindex(p.shape): old = p.data[i] p.data[i] = old + eps with zc.no_grad(): up = (module(xt).data * w).sum() p.data[i] = old - eps with zc.no_grad(): down = (module(xt).data * w).sum() p.data[i] = old num[i] = (up - down) / (2 * eps) err = np.abs(num - p.grad).max() / max(1.0, np.abs(num).max()) assert err < 1e-5, f"参数梯度误差 {err:.2e}" # ---------------- 算子 ---------------- def test_ops(): grad_check(lambda a, b: a + b, (3, 4), (4,)) grad_check(lambda a, b: a - b, (3, 4), (3, 1)) grad_check(lambda a, b: a * b, (2, 3), (2, 3)) grad_check(lambda a, b: a / b, (2, 3), (3,), positive=True) grad_check(lambda a: a ** 3, (4,)) grad_check(lambda a: a.sqrt(), (4,), positive=True) grad_check(lambda a: -a + 2.0 - a * 3.0, (4,)) grad_check(lambda a: 1.0 / a, (4,), positive=True) grad_check(lambda a: a.exp(), (3, 2)) grad_check(lambda a: a.log(), (3, 2), positive=True) grad_check(lambda a: a.tanh(), (5,)) grad_check(lambda a: a.relu(), (5,)) grad_check(lambda a: a.masked_fill(np.eye(3, dtype=bool), -5.0), (3, 3)) def test_reductions_and_shapes(): grad_check(lambda a: a.sum(), (3, 4)) grad_check(lambda a: a.sum(axis=1), (3, 4)) grad_check(lambda a: a.mean(axis=0, keepdims=True), (3, 4)) grad_check(lambda a: a.max(axis=-1), (3, 4)) grad_check(lambda a: a.reshape(4, 3), (3, 4)) grad_check(lambda a: a.transpose(2, 0, 1), (2, 3, 4)) grad_check(lambda a: a[1:, ::2], (3, 4)) grad_check(lambda a: a[np.array([0, 2, 0])], (3, 4)) # 重复索引 grad_check(lambda a, b: zc.concat([a, b, a], axis=1), (2, 3), (2, 1)) def test_matmul(): grad_check(lambda a, b: a @ b, (3, 4), (4, 5)) grad_check(lambda a, b: a @ b, (2, 3, 4), (4, 5)) # 批量 @ 共享权重 grad_check(lambda a, b: a @ b, (2, 2, 3, 4), (2, 2, 4, 3)) def test_functional(): grad_check(lambda a: F.softmax(a, axis=-1), (3, 5)) grad_check(lambda a: F.log_softmax(a, axis=-1), (3, 5)) grad_check(lambda a: F.gelu(a), (6,)) grad_check(lambda a: F.sigmoid(a), (6,)) grad_check(lambda a: F.cross_entropy(a, np.array([1, 0, 4])), (3, 5)) # 和直接的数学定义对比数值 x = rng.standard_normal((4, 6)) t = np.array([0, 5, 2, 2]) ref = -np.log(np.exp(x) / np.exp(x).sum(1, keepdims=True))[np.arange(4), t].mean() assert abs(F.cross_entropy(zc.Tensor(x), t).item() - ref) < 1e-10 def test_layers(): module_grad_check(nn.Linear(4, 3), rng.standard_normal((2, 4))) module_grad_check(nn.LayerNorm(5), rng.standard_normal((3, 5))) module_grad_check(nn.Sequential(nn.Linear(3, 4), nn.GELU(), nn.Linear(4, 2)), rng.standard_normal((2, 3))) module_grad_check(nn.CausalSelfAttention(8, 2), rng.standard_normal((2, 4, 8))) emb = nn.Embedding(5, 3) out = emb(np.array([[1, 1, 4]])) out.sum().backward() assert emb.weight.grad[1].tolist() == [2, 2, 2] and emb.weight.grad[0].tolist() == [0, 0, 0] def test_attention_is_causal(): attn = nn.CausalSelfAttention(8, 2).eval() x = rng.standard_normal((1, 5, 8)).astype(np.float32) y1 = attn(zc.Tensor(x)).data x[0, 4] += 10 # 改最后一个位置，前面的输出不应变化 y2 = attn(zc.Tensor(x)).data assert np.allclose(y1[0, :4], y2[0, :4]) and not np.allclose(y1[0, 4], y2[0, 4]) def test_count_flops(): with zc.count_flops() as c: zc.randn(2, 3, 4) @ zc.randn(4, 5) assert c.flops == 2 * (2 * 3 * 5) * 4 def test_default_dtype(): zc.set_default_dtype(np.float32) try: x = zc.Tensor(np.ones(3)) # float64 数组进来也变成 float32 assert x.dtype == np.float32 and (x * 0.5).dtype == np.float32 assert nn.Linear(2, 2).weight.dtype == np.float32 assert zc.Tensor(np.arange(3)).dtype.kind == "i" # 整数下标不动 assert np.isfinite(F.sigmoid(zc.Tensor([-1000.0, 0.0, 1000.0])).data).all() finally: zc.set_default_dtype(np.float64) def test_state_dict_roundtrip(tmp="/tmp"): path = os.path.join(os.environ.get("TMPDIR", tmp), "zerocg_test.npz") a = nn.Sequential(nn.Linear(3, 4), nn.LayerNorm(4)) a.save(path) b = nn.Sequential(nn.Linear(3, 4), nn.LayerNorm(4)).load(path) x = zc.randn(2, 3) assert np.allclose(a(x).data, b(x).data) assert [n for n, _ in a.named_parameters()] == ["layers.0.weight", "layers.0.bias", "layers.1.weight", "layers.1.bias"] def test_optimizer_resume(): """训练 10 步 == 训练 5 步、存档、新建模型和优化器、读档、再训练 5 步。""" X, y = rng.standard_normal((8, 3)), np.array([0, 1, 1, 0, 1, 0, 0, 1]) def make(): np.random.seed(7) m = nn.Sequential(nn.Linear(3, 4), nn.Tanh(), nn.Linear(4, 2)) return m, zc.optim.AdamW(m.parameters(), lr=0.05, weight_decay=0.01) def train(m, opt, n): for _ in range(n): loss = F.cross_entropy(m(zc.Tensor(X)), y) opt.zero_grad() loss.backward() opt.step() a, opt_a = make() train(a, opt_a, 10) b, opt_b = make() train(b, opt_b, 5) model_state, opt_state = b.state_dict(), opt_b.state_dict() c, opt_c = make() c.load_state_dict(model_state) opt_c.load_state_dict(opt_state) train(c, opt_c, 5) assert all(np.array_equal(p, q) for p, q in zip(a.state_dict().values(), c.state_dict().values())) # ---------------- 端到端学习 ---------------- def test_learns_xor(): np.random.seed(0) X = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=np.float32) y = np.array([0, 1, 1, 0]) model = nn.Sequential(nn.Linear(2, 16), nn.Tanh(), nn.Linear(16, 2)) opt = zc.optim.AdamW(model.parameters(), lr=0.05) for _ in range(300): loss = F.cross_entropy(model(zc.Tensor(X)), y) opt.zero_grad() loss.backward() opt.step() assert (model(zc.Tensor(X)).data.argmax(1) == y).all(), f"XOR 没学会, loss={loss.item()}" def test_attention_memorizes_sequence(): """一层注意力 + 嵌入，记住一个短序列的下一个字符。""" np.random.seed(0) seq = np.array([0, 1, 2, 3, 2, 1, 0, 3]) x, y = seq[None, :-1], seq[1:] class Tiny(nn.Module): def __init__(self): self.tok = nn.Embedding(4, 16) self.pos = nn.Embedding(8, 16) self.attn = nn.CausalSelfAttention(16, 2) self.ln = nn.LayerNorm(16) self.head = nn.Linear(16, 4) def forward(self, idx): T = idx.shape[1] h = self.tok(idx) + self.pos(np.arange(T)) h = h + self.attn(self.ln(h)) return self.head(h) model = Tiny() opt = zc.optim.AdamW(model.parameters(), lr=0.02) for _ in range(200): loss = F.cross_entropy(model(x).reshape(-1, 4), y) opt.zero_grad() loss.backward() zc.optim.clip_grad_norm(model.parameters(), 1.0) opt.step() assert loss.item() < 0.05, f"loss={loss.item()}" if __name__ == "__main__": tests = [(k, v) for k, v in list(globals().items()) if k.startswith("test_")] failed = 0 for name, fn in tests: try: fn() print(f" ok {name}") except Exception as e: # noqa: BLE001 failed += 1 print(f" FAIL {name}: {type(e).__name__}: {e}") print(f"\n{len(tests) - failed}/{len(tests)} 通过") sys.exit(1 if failed else 0)