ZeroPoint v0 research prototype: new from-scratch NumPy framework, memory-slot architecture, audited Transformer comparison
5c6c847 verified | """常用函数。softmax / log_softmax / gelu 写成融合算子（手写反向），更快也更数值稳定。""" | |
| import numpy as np | |
| from .tensor import Tensor | |
| def relu(x): | |
| return x.relu() | |
| def tanh(x): | |
| return x.tanh() | |
| def sigmoid_np(a): | |
| """数值稳定：只对 -|a| 求 exp，永远不会溢出。""" | |
| e = np.exp(-np.abs(a)) | |
| return np.where(a >= 0, 1 / (1 + e), e / (1 + e)) | |
| def sigmoid(x): | |
| out = sigmoid_np(x.data) | |
| return Tensor._op(out, (x,), lambda g: (g * out * (1 - out),)) | |
| def gelu(x): | |
| """GELU（tanh 近似，GPT-2 同款）。""" | |
| c = np.sqrt(2 / np.pi) | |
| a = x.data | |
| inner = c * (a + 0.044715 * a ** 3) | |
| t = np.tanh(inner) | |
| out = 0.5 * a * (1 + t) | |
| def backward(g): | |
| dinner = c * (1 + 3 * 0.044715 * a ** 2) | |
| return (g * (0.5 * (1 + t) + 0.5 * a * (1 - t ** 2) * dinner),) | |
| return Tensor._op(out.astype(a.dtype), (x,), backward) | |
| def softmax(x, axis=-1): | |
| e = np.exp(x.data - x.data.max(axis=axis, keepdims=True)) | |
| s = e / e.sum(axis=axis, keepdims=True) | |
| return Tensor._op(s, (x,), lambda g: (s * (g - (g * s).sum(axis=axis, keepdims=True)),)) | |
| def log_softmax(x, axis=-1): | |
| shifted = x.data - x.data.max(axis=axis, keepdims=True) | |
| out = shifted - np.log(np.exp(shifted).sum(axis=axis, keepdims=True)) | |
| s = np.exp(out) | |
| return Tensor._op(out, (x,), lambda g: (g - s * g.sum(axis=axis, keepdims=True),)) | |
| def cross_entropy(logits, targets): | |
| """logits: (N, C)；targets: (N,) 整数类别。返回平均负对数似然。""" | |
| targets = targets.data if isinstance(targets, Tensor) else np.asarray(targets) | |
| logp = log_softmax(logits, axis=-1) | |
| return -logp[np.arange(len(targets)), targets].mean() | |
| def mse_loss(pred, target): | |
| return ((pred - target) ** 2).mean() | |
| def dropout(x, p, training=True): | |
| if not training or p == 0: | |
| return x | |
| mask = (np.random.rand(*x.shape) >= p).astype(x.dtype) / (1 - p) | |
| return x * Tensor(mask) | |