ZeroPoint v0 research prototype: new from-scratch NumPy framework, memory-slot architecture, audited Transformer comparison
5c6c847 verified | """神经网络层。Module 会自动收集属性里的 Parameter / 子 Module。""" | |
| import numpy as np | |
| from . import functional as F | |
| from .tensor import Tensor, get_default_dtype | |
| class Parameter(Tensor): | |
| def __init__(self, data): | |
| super().__init__(np.asarray(data, dtype=get_default_dtype()), requires_grad=True) | |
| class Module: | |
| training = True | |
| def forward(self, *args, **kwargs): | |
| raise NotImplementedError | |
| def __call__(self, *args, **kwargs): | |
| return self.forward(*args, **kwargs) | |
| def named_parameters(self, prefix=""): | |
| for name, v in vars(self).items(): | |
| items = [(name, v)] if not isinstance(v, (list, tuple)) else [(f"{name}.{i}", x) for i, x in enumerate(v)] | |
| for n, x in items: | |
| if isinstance(x, Parameter): | |
| yield prefix + n, x | |
| elif isinstance(x, Module): | |
| yield from x.named_parameters(prefix + n + ".") | |
| def parameters(self): | |
| return [p for _, p in self.named_parameters()] | |
| def modules(self): | |
| yield self | |
| for v in vars(self).values(): | |
| for x in (v if isinstance(v, (list, tuple)) else [v]): | |
| if isinstance(x, Module): | |
| yield from x.modules() | |
| def train(self, mode=True): | |
| for m in self.modules(): | |
| m.training = mode | |
| return self | |
| def eval(self): | |
| return self.train(False) | |
| def zero_grad(self): | |
| for p in self.parameters(): | |
| p.grad = None | |
| def num_parameters(self): | |
| return sum(p.data.size for p in self.parameters()) | |
| # ---------- 保存 / 加载 ---------- | |
| def state_dict(self): | |
| return {n: p.data.copy() for n, p in self.named_parameters()} | |
| def load_state_dict(self, state): | |
| for n, p in self.named_parameters(): | |
| assert state[n].shape == p.shape, f"{n} 形状不匹配: {state[n].shape} vs {p.shape}" | |
| p.data = state[n].astype(p.dtype) | |
| def save(self, path): | |
| np.savez(path, **self.state_dict()) | |
| def load(self, path): | |
| with np.load(path) as f: | |
| self.load_state_dict(dict(f)) | |
| return self | |
| class Linear(Module): | |
| def __init__(self, in_features, out_features, bias=True): | |
| bound = 1 / np.sqrt(in_features) | |
| self.weight = Parameter(np.random.uniform(-bound, bound, (in_features, out_features))) | |
| self.bias = Parameter(np.zeros(out_features)) if bias else None | |
| def forward(self, x): | |
| y = x @ self.weight | |
| return y + self.bias if self.bias is not None else y | |
| class Embedding(Module): | |
| def __init__(self, num_embeddings, dim): | |
| self.weight = Parameter(np.random.randn(num_embeddings, dim) * 0.02) | |
| def forward(self, idx): | |
| return self.weight[idx.data if isinstance(idx, Tensor) else np.asarray(idx)] | |
| class LayerNorm(Module): | |
| def __init__(self, dim, eps=1e-5): | |
| self.eps = eps | |
| self.weight = Parameter(np.ones(dim)) | |
| self.bias = Parameter(np.zeros(dim)) | |
| def forward(self, x): | |
| mu = x.mean(axis=-1, keepdims=True) | |
| xc = x - mu | |
| var = (xc * xc).mean(axis=-1, keepdims=True) | |
| return xc / (var + self.eps).sqrt() * self.weight + self.bias | |
| class Dropout(Module): | |
| def __init__(self, p=0.1): | |
| self.p = p | |
| def forward(self, x): | |
| return F.dropout(x, self.p, self.training) | |
| class ReLU(Module): | |
| def forward(self, x): | |
| return x.relu() | |
| class Tanh(Module): | |
| def forward(self, x): | |
| return x.tanh() | |
| class GELU(Module): | |
| def forward(self, x): | |
| return F.gelu(x) | |
| class Sequential(Module): | |
| def __init__(self, *layers): | |
| self.layers = list(layers) | |
| def forward(self, x): | |
| for layer in self.layers: | |
| x = layer(x) | |
| return x | |
| def __getitem__(self, i): | |
| return self.layers[i] | |
| class CausalSelfAttention(Module): | |
| """多头因果自注意力：位置 t 只能看到 <= t 的位置。输入/输出形状 (B, T, C)。""" | |
| def __init__(self, n_embd, n_head, dropout=0.0): | |
| assert n_embd % n_head == 0 | |
| self.n_head = n_head | |
| self.qkv = Linear(n_embd, 3 * n_embd) | |
| self.proj = Linear(n_embd, n_embd) | |
| self.attn_drop = Dropout(dropout) | |
| self.resid_drop = Dropout(dropout) | |
| def forward(self, x): | |
| B, T, C = x.shape | |
| H, D = self.n_head, C // self.n_head | |
| qkv = self.qkv(x).reshape(B, T, 3, H, D).transpose(2, 0, 3, 1, 4) # (3, B, H, T, D) | |
| q, k, v = qkv[0], qkv[1], qkv[2] | |
| att = (q @ k.transpose(0, 1, 3, 2)) * (1.0 / np.sqrt(D)) # (B, H, T, T) | |
| att = att.masked_fill(np.triu(np.ones((T, T), dtype=bool), 1), -1e9) | |
| att = self.attn_drop(F.softmax(att, axis=-1)) | |
| y = (att @ v).transpose(0, 2, 1, 3).reshape(B, T, C) | |
| return self.resid_drop(self.proj(y)) | |