ZeroPoint v0 research prototype: new from-scratch NumPy framework, memory-slot architecture, audited Transformer comparison
5c6c847 verified | """Tensor + 反向模式自动求导（只依赖 numpy）。 | |
| 每个运算产出一个新 Tensor，并记录: | |
| _parents : 参与运算的输入 Tensor | |
| _backward : 函数 g -> (对每个 parent 的梯度, ...) | |
| backward() 先拓扑排序，再从输出往回逐个调用 _backward，把梯度累加到 .grad。 | |
| """ | |
| import numpy as np | |
| _grad_enabled = True | |
| _default_dtype = np.dtype(np.float32) | |
| def set_default_dtype(dtype): | |
| """所有浮点数据进入框架时统一转成这个类型（默认 float32）。 | |
| 只有测试里的数值梯度检查需要切到 float64。整数数组（下标）不受影响。""" | |
| global _default_dtype | |
| _default_dtype = np.dtype(dtype) | |
| def get_default_dtype(): | |
| return _default_dtype | |
| _flops = None # count_flops() 打开时是一个 [计数]，否则 None | |
| class count_flops: | |
| """with count_flops() as c: ... 之后 c.flops = 期间所有矩阵乘的浮点运算次数（乘加算 2 次）。 | |
| 只统计矩阵乘（模型里绝大部分计算量），不统计逐元素运算；不受机器负载影响，比计时可靠。""" | |
| def __enter__(self): | |
| global _flops | |
| self._prev, _flops = _flops, [0] | |
| return self | |
| def __exit__(self, *exc): | |
| global _flops | |
| self.flops = _flops[0] | |
| _flops = self._prev | |
| class no_grad: | |
| """with no_grad(): 内部的运算不建计算图（推理/评估时用，省内存）。""" | |
| def __enter__(self): | |
| global _grad_enabled | |
| self._prev, _grad_enabled = _grad_enabled, False | |
| def __exit__(self, *exc): | |
| global _grad_enabled | |
| _grad_enabled = self._prev | |
| def _unbroadcast(grad, shape): | |
| """把广播后的梯度求和还原成原始形状。""" | |
| while grad.ndim > len(shape): | |
| grad = grad.sum(axis=0) | |
| for i, s in enumerate(shape): | |
| if s == 1 and grad.shape[i] != 1: | |
| grad = grad.sum(axis=i, keepdims=True) | |
| return grad | |
| def _to_array(x): | |
| if isinstance(x, Tensor): | |
| return x.data | |
| arr = np.asarray(x) | |
| # 浮点一律转成默认类型，避免 float32 和 float64 混在一起被 numpy 悄悄升成 float64 | |
| if arr.dtype.kind == "f" and arr.dtype != _default_dtype: | |
| arr = arr.astype(_default_dtype) | |
| return arr | |
| class Tensor: | |
| __array_priority__ = 100 # 让 ndarray + Tensor 走 Tensor.__radd__ | |
| def __init__(self, data, requires_grad=False): | |
| self.data = _to_array(data) | |
| self.requires_grad = requires_grad | |
| self.grad = None | |
| self._parents = () | |
| self._backward = None | |
| # ---------- 基础属性 ---------- | |
| def shape(self): | |
| return self.data.shape | |
| def ndim(self): | |
| return self.data.ndim | |
| def dtype(self): | |
| return self.data.dtype | |
| def __len__(self): | |
| return len(self.data) | |
| def __repr__(self): | |
| return f"Tensor({self.data}, requires_grad={self.requires_grad})" | |
| def numpy(self): | |
| return self.data | |
| def item(self): | |
| return self.data.item() | |
| def detach(self): | |
| return Tensor(self.data) | |
| # ---------- 建图 ---------- | |
| def _op(data, parents, backward): | |
| """创建运算结果。backward(g) 返回与 parents 一一对应的梯度（可为 None）。""" | |
| out = Tensor(data) | |
| if _grad_enabled and any(p.requires_grad for p in parents): | |
| out.requires_grad = True | |
| out._parents = parents | |
| out._backward = backward | |
| return out | |
| def backward(self, grad=None): | |
| if grad is None: | |
| if self.data.size != 1: | |
| raise RuntimeError("非标量输出调用 backward() 需要显式传入 grad") | |
| grad = np.ones_like(self.data) | |
| # 迭代式拓扑排序（避免深图递归爆栈） | |
| topo, visited, stack = [], set(), [(self, False)] | |
| while stack: | |
| node, expanded = stack.pop() | |
| if expanded: | |
| topo.append(node) | |
| continue | |
| if id(node) in visited: | |
| continue | |
| visited.add(id(node)) | |
| stack.append((node, True)) | |
| for p in node._parents: | |
| if p.requires_grad and id(p) not in visited: | |
| stack.append((p, False)) | |
| self.grad = grad if self.grad is None else self.grad + grad | |
| for node in reversed(topo): | |
| if node._backward is None or node.grad is None: | |
| continue | |
| for p, g in zip(node._parents, node._backward(node.grad)): | |
| if g is not None and p.requires_grad: | |
| p.grad = g if p.grad is None else p.grad + g | |
| # 释放计算图，中间节点的梯度也不再需要 | |
| node._parents, node._backward, node.grad = (), None, None | |
| # ---------- 逐元素运算 ---------- | |
| def __add__(self, other): | |
| other = other if isinstance(other, Tensor) else Tensor(other) | |
| a, b = self, other | |
| return Tensor._op(a.data + b.data, (a, b), | |
| lambda g: (_unbroadcast(g, a.shape), _unbroadcast(g, b.shape))) | |
| def __mul__(self, other): | |
| other = other if isinstance(other, Tensor) else Tensor(other) | |
| a, b = self, other | |
| return Tensor._op(a.data * b.data, (a, b), | |
| lambda g: (_unbroadcast(g * b.data, a.shape), _unbroadcast(g * a.data, b.shape))) | |
| def __truediv__(self, other): | |
| other = other if isinstance(other, Tensor) else Tensor(other) | |
| a, b = self, other | |
| return Tensor._op(a.data / b.data, (a, b), | |
| lambda g: (_unbroadcast(g / b.data, a.shape), | |
| _unbroadcast(-g * a.data / (b.data ** 2), b.shape))) | |
| def __pow__(self, p): | |
| assert isinstance(p, (int, float)), "只支持标量指数" | |
| a = self | |
| return Tensor._op(a.data ** p, (a,), lambda g: (g * p * a.data ** (p - 1),)) | |
| def __neg__(self): | |
| return Tensor._op(-self.data, (self,), lambda g: (-g,)) | |
| def __sub__(self, other): | |
| return self + (-other if isinstance(other, Tensor) else Tensor(-_to_array(other))) | |
| def __radd__(self, other): | |
| return self + other | |
| def __rmul__(self, other): | |
| return self * other | |
| def __rsub__(self, other): | |
| return (-self) + other | |
| def __rtruediv__(self, other): | |
| return Tensor(other) / self | |
| def exp(self): | |
| out = np.exp(self.data) | |
| return Tensor._op(out, (self,), lambda g: (g * out,)) | |
| def log(self): | |
| a = self | |
| return Tensor._op(np.log(a.data), (a,), lambda g: (g / a.data,)) | |
| def sqrt(self): | |
| return self ** 0.5 | |
| def tanh(self): | |
| out = np.tanh(self.data) | |
| return Tensor._op(out, (self,), lambda g: (g * (1 - out ** 2),)) | |
| def relu(self): | |
| mask = self.data > 0 | |
| return Tensor._op(self.data * mask, (self,), lambda g: (g * mask,)) | |
| def masked_fill(self, mask, value): | |
| mask = _to_array(mask).astype(bool) | |
| return Tensor._op(np.where(mask, value, self.data).astype(self.dtype), (self,), | |
| lambda g: (np.where(mask, 0, g),)) | |
| # ---------- 归约 ---------- | |
| def sum(self, axis=None, keepdims=False): | |
| a = self | |
| def backward(g): | |
| if axis is not None and not keepdims: | |
| g = np.expand_dims(g, axis) | |
| return (np.broadcast_to(g, a.shape).copy(),) | |
| return Tensor._op(a.data.sum(axis=axis, keepdims=keepdims), (a,), backward) | |
| def mean(self, axis=None, keepdims=False): | |
| n = self.data.size if axis is None else np.prod([self.shape[i] for i in np.atleast_1d(axis)]) | |
| return self.sum(axis, keepdims) * (1.0 / n) | |
| def max(self, axis=None, keepdims=False): | |
| a = self | |
| out = a.data.max(axis=axis, keepdims=True) | |
| def backward(g): | |
| if axis is not None and not keepdims: | |
| g = np.expand_dims(g, axis) | |
| mask = a.data == out | |
| return (mask * g / mask.sum(axis=axis, keepdims=True),) # 并列最大值平分梯度 | |
| return Tensor._op(out if keepdims else out.squeeze(axis), (a,), backward) | |
| # ---------- 形状 ---------- | |
| def reshape(self, *shape): | |
| a = self | |
| return Tensor._op(a.data.reshape(*shape), (a,), lambda g: (g.reshape(a.shape),)) | |
| def transpose(self, *axes): | |
| axes = axes or tuple(reversed(range(self.ndim))) | |
| inv = np.argsort(axes) | |
| return Tensor._op(self.data.transpose(axes), (self,), lambda g: (g.transpose(inv),)) | |
| def T(self): | |
| return self.transpose() | |
| def __getitem__(self, idx): | |
| a = self | |
| if isinstance(idx, Tensor): | |
| idx = idx.data | |
| parts = idx if isinstance(idx, tuple) else (idx,) | |
| fancy = any(isinstance(i, (np.ndarray, list)) for i in parts) | |
| def backward(g): | |
| grad = np.zeros_like(a.data) | |
| if fancy: | |
| np.add.at(grad, idx, g) # 数组下标可能重复（Embedding 查表），必须累加 | |
| else: | |
| grad[idx] = g # 普通切片不会重复，直接写，比 add.at 快得多 | |
| return (grad,) | |
| return Tensor._op(a.data[idx], (a,), backward) | |
| # ---------- 矩阵乘 ---------- | |
| def __matmul__(self, other): | |
| other = other if isinstance(other, Tensor) else Tensor(other) | |
| a, b = self, other | |
| assert a.ndim >= 2 and b.ndim >= 2, "matmul 需要至少 2 维" | |
| out = a.data @ b.data | |
| if _flops is not None: | |
| _flops[0] += 2 * out.size * a.shape[-1] | |
| def backward(g): | |
| ga = g @ np.swapaxes(b.data, -1, -2) | |
| gb = np.swapaxes(a.data, -1, -2) @ g | |
| return _unbroadcast(ga, a.shape), _unbroadcast(gb, b.shape) | |
| return Tensor._op(out, (a, b), backward) | |
| def concat(tensors, axis=0): | |
| """沿 axis 拼接。反向时把梯度按原来的大小切回去。""" | |
| tensors = tuple(t if isinstance(t, Tensor) else Tensor(t) for t in tensors) | |
| splits = np.cumsum([t.shape[axis] for t in tensors])[:-1] | |
| return Tensor._op(np.concatenate([t.data for t in tensors], axis=axis), tensors, | |
| lambda g: tuple(np.split(g, splits, axis=axis))) | |
| # ---------- 构造函数 ---------- | |
| def tensor(data, requires_grad=False): | |
| return Tensor(data, requires_grad) | |
| def zeros(*shape, requires_grad=False): | |
| return Tensor(np.zeros(shape, dtype=_default_dtype), requires_grad) | |
| def ones(*shape, requires_grad=False): | |
| return Tensor(np.ones(shape, dtype=_default_dtype), requires_grad) | |
| def randn(*shape, requires_grad=False): | |
| return Tensor(np.random.randn(*shape).astype(_default_dtype), requires_grad) | |