File size: 4,882 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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | """神经网络层。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))
|