File size: 1,976 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 | """常用函数。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)
|