257 lines
8.7 KiB
Python
257 lines
8.7 KiB
Python
|
"""
|
|||
|
|
训练轻量 Transformer 记忆模型，导出 weights.json（NumPy 推理）。
|
|||
|
|
|
|||
|
|
用法:
|
|||
|
|
cd backend && source venv/bin/activate
|
|||
|
|
pip install torch numpy # 训练仅需本机安装
|
|||
|
|
python -m scripts.train_memory_transformer
|
|||
|
|
python -m scripts.train_memory_transformer --epochs 30 --from-db
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import argparse
|
|||
|
|
import math
|
|||
|
|
import sys
|
|||
|
|
from datetime import datetime, timezone
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
import numpy as np
|
|||
|
|
|
|||
|
|
BACKEND = Path(__file__).resolve().parents[1]
|
|||
|
|
sys.path.insert(0, str(BACKEND))
|
|||
|
|
|
|||
|
|
from database import SessionLocal # noqa: E402
|
|||
|
|
from models import QuizRecord, Word # noqa: E402
|
|||
|
|
from services.memory_transformer.encoding import ( # noqa: E402
|
|||
|
|
FEATURE_DIM,
|
|||
|
|
MAX_SEQ_LEN,
|
|||
|
|
build_sequence_matrix,
|
|||
|
|
)
|
|||
|
|
from services.memory_transformer.network import MiniTransformer, init_random_weights # noqa: E402
|
|||
|
|
from services.memory_transformer.service import ( # noqa: E402
|
|||
|
|
DEFAULT_WEIGHTS_PATH,
|
|||
|
|
WEIGHTS_PATH,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
import torch
|
|||
|
|
import torch.nn as nn
|
|||
|
|
except ImportError:
|
|||
|
|
torch = None # type: ignore
|
|||
|
|
nn = None # type: ignore
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _build_torch_model():
|
|||
|
|
class TorchMiniTransformer(nn.Module):
|
|||
|
|
def __init__(self, d_model=48, n_heads=2, d_ff=96, n_layers=2):
|
|||
|
|
super().__init__()
|
|||
|
|
self.d_model = d_model
|
|||
|
|
self.n_heads = n_heads
|
|||
|
|
self.n_layers = n_layers
|
|||
|
|
self.in_proj = nn.Linear(FEATURE_DIM, d_model)
|
|||
|
|
self.ln_in = nn.LayerNorm(d_model)
|
|||
|
|
self.layers = nn.ModuleList(
|
|||
|
|
[
|
|||
|
|
nn.TransformerEncoderLayer(
|
|||
|
|
d_model=d_model,
|
|||
|
|
nhead=n_heads,
|
|||
|
|
dim_feedforward=d_ff,
|
|||
|
|
batch_first=True,
|
|||
|
|
dropout=0.1,
|
|||
|
|
)
|
|||
|
|
for _ in range(n_layers)
|
|||
|
|
]
|
|||
|
|
)
|
|||
|
|
self.head = nn.Linear(d_model, 1)
|
|||
|
|
|
|||
|
|
def forward(self, x, valid_lens):
|
|||
|
|
h = self.ln_in(self.in_proj(x))
|
|||
|
|
key_padding = torch.zeros(
|
|||
|
|
x.size(0), MAX_SEQ_LEN, dtype=torch.bool, device=x.device
|
|||
|
|
)
|
|||
|
|
for i, vl in enumerate(valid_lens):
|
|||
|
|
if vl < MAX_SEQ_LEN:
|
|||
|
|
key_padding[i, vl:] = True
|
|||
|
|
for layer in self.layers:
|
|||
|
|
h = layer(h, src_key_padding_mask=key_padding)
|
|||
|
|
cls = h[:, 0, :]
|
|||
|
|
return self.head(cls).squeeze(-1)
|
|||
|
|
|
|||
|
|
return TorchMiniTransformer()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def export_torch_to_numpy(torch_model) -> MiniTransformer:
|
|||
|
|
"""将 PyTorch 权重映射到 NumPy MiniTransformer 命名。"""
|
|||
|
|
m = MiniTransformer(
|
|||
|
|
d_model=torch_model.d_model,
|
|||
|
|
n_heads=torch_model.n_heads,
|
|||
|
|
d_ff=96,
|
|||
|
|
n_layers=torch_model.n_layers,
|
|||
|
|
)
|
|||
|
|
w = m.weights
|
|||
|
|
w["in_proj"] = torch_model.in_proj.weight.detach().cpu().numpy().T
|
|||
|
|
w["in_bias"] = torch_model.in_proj.bias.detach().cpu().numpy()
|
|||
|
|
w["ln_in_g"] = torch_model.ln_in.weight.detach().cpu().numpy()
|
|||
|
|
w["ln_in_b"] = torch_model.ln_in.bias.detach().cpu().numpy()
|
|||
|
|
|
|||
|
|
for li, layer in enumerate(torch_model.layers):
|
|||
|
|
attn = layer.self_attn
|
|||
|
|
d = m.d_model
|
|||
|
|
# PyTorch MHA: in_proj_weight stacks Q,K,V
|
|||
|
|
in_w = attn.in_proj_weight.detach().cpu().numpy()
|
|||
|
|
Wq, Wk, Wv = in_w[:d], in_w[d : 2 * d], in_w[2 * d :]
|
|||
|
|
w[f"L{li}.Wq"] = Wq.T
|
|||
|
|
w[f"L{li}.Wk"] = Wk.T
|
|||
|
|
w[f"L{li}.Wv"] = Wv.T
|
|||
|
|
in_b = attn.in_proj_bias.detach().cpu().numpy()
|
|||
|
|
w[f"L{li}.Bq"] = in_b[:d]
|
|||
|
|
w[f"L{li}.Bk"] = in_b[d : 2 * d]
|
|||
|
|
w[f"L{li}.Bv"] = in_b[2 * d :]
|
|||
|
|
w[f"L{li}.Wo"] = attn.out_proj.weight.detach().cpu().numpy().T
|
|||
|
|
w[f"L{li}.Bo"] = attn.out_proj.bias.detach().cpu().numpy()
|
|||
|
|
w[f"L{li}.ln1_g"] = layer.norm1.weight.detach().cpu().numpy()
|
|||
|
|
w[f"L{li}.ln1_b"] = layer.norm1.bias.detach().cpu().numpy()
|
|||
|
|
w[f"L{li}.W1"] = layer.linear1.weight.detach().cpu().numpy().T
|
|||
|
|
w[f"L{li}.b1"] = layer.linear1.bias.detach().cpu().numpy()
|
|||
|
|
w[f"L{li}.W2"] = layer.linear2.weight.detach().cpu().numpy().T
|
|||
|
|
w[f"L{li}.b2"] = layer.linear2.bias.detach().cpu().numpy()
|
|||
|
|
w[f"L{li}.ln2_g"] = layer.norm2.weight.detach().cpu().numpy()
|
|||
|
|
w[f"L{li}.ln2_b"] = layer.norm2.bias.detach().cpu().numpy()
|
|||
|
|
|
|||
|
|
w["head_w"] = torch_model.head.weight.detach().cpu().numpy()[0]
|
|||
|
|
w["head_b"] = torch_model.head.bias.detach().cpu().numpy()[0]
|
|||
|
|
return m
|
|||
|
|
|
|||
|
|
|
|||
|
|
def load_samples_from_db(limit: int = 50000) -> tuple[list, list]:
|
|||
|
|
db = SessionLocal()
|
|||
|
|
try:
|
|||
|
|
records = (
|
|||
|
|
db.query(QuizRecord)
|
|||
|
|
.order_by(QuizRecord.created_at.asc())
|
|||
|
|
.limit(limit)
|
|||
|
|
.all()
|
|||
|
|
)
|
|||
|
|
word_cache: dict[int, Word] = {}
|
|||
|
|
xs, ys = [], []
|
|||
|
|
for r in records:
|
|||
|
|
if r.word_id not in word_cache:
|
|||
|
|
word_cache[r.word_id] = db.query(Word).filter(Word.id == r.word_id).first()
|
|||
|
|
word = word_cache.get(r.word_id)
|
|||
|
|
if not word:
|
|||
|
|
continue
|
|||
|
|
hist = (
|
|||
|
|
db.query(QuizRecord)
|
|||
|
|
.filter(QuizRecord.word_id == r.word_id, QuizRecord.created_at < r.created_at)
|
|||
|
|
.order_by(QuizRecord.created_at.asc())
|
|||
|
|
.all()
|
|||
|
|
)
|
|||
|
|
seq, vl = build_sequence_matrix(word, hist, parse_iso(r.created_at))
|
|||
|
|
xs.append(seq)
|
|||
|
|
ys.append(1.0 if r.is_correct else 0.0)
|
|||
|
|
return xs, ys
|
|||
|
|
finally:
|
|||
|
|
db.close()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def parse_iso(s: str) -> datetime:
|
|||
|
|
s = s.replace("Z", "+00:00")
|
|||
|
|
try:
|
|||
|
|
return datetime.fromisoformat(s)
|
|||
|
|
except ValueError:
|
|||
|
|
return datetime.strptime(s[:19], "%Y-%m-%dT%H:%M:%S").replace(tzinfo=timezone.utc)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def synthetic_samples(n: int = 4000) -> tuple[list, list]:
|
|||
|
|
rng = np.random.default_rng(0)
|
|||
|
|
xs, ys = [], []
|
|||
|
|
for _ in range(n):
|
|||
|
|
vl = int(rng.integers(2, 12))
|
|||
|
|
seq = rng.normal(0, 0.5, (MAX_SEQ_LEN, FEATURE_DIM)).tolist()
|
|||
|
|
seq[0][0] = rng.uniform(0, 1)
|
|||
|
|
seq[0][1] = rng.uniform(0, 1)
|
|||
|
|
last_ok = seq[-1][0] if vl > 1 else 0.5
|
|||
|
|
y = 1.0 if rng.random() < 0.4 + 0.4 * last_ok + 0.1 * seq[0][0] else 0.0
|
|||
|
|
xs.append(seq)
|
|||
|
|
ys.append(y)
|
|||
|
|
return xs, ys
|
|||
|
|
|
|||
|
|
|
|||
|
|
def train_and_export(
|
|||
|
|
epochs: int = 20,
|
|||
|
|
batch_size: int = 64,
|
|||
|
|
from_db: bool = False,
|
|||
|
|
out_path: Path = WEIGHTS_PATH,
|
|||
|
|
) -> None:
|
|||
|
|
if torch is None:
|
|||
|
|
print("请先安装 PyTorch: pip install torch")
|
|||
|
|
print("将写入随机初始化 default_weights 供推理回退…")
|
|||
|
|
m = MiniTransformer()
|
|||
|
|
m.weights = init_random_weights()
|
|||
|
|
DEFAULT_WEIGHTS_PATH.parent.mkdir(parents=True, exist_ok=True)
|
|||
|
|
m.save_json(DEFAULT_WEIGHTS_PATH)
|
|||
|
|
print(f"已保存 {DEFAULT_WEIGHTS_PATH}")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
xs, ys = load_samples_from_db() if from_db else synthetic_samples()
|
|||
|
|
if len(xs) < 32:
|
|||
|
|
print("样本不足，使用合成数据")
|
|||
|
|
xs, ys = synthetic_samples()
|
|||
|
|
|
|||
|
|
X = torch.tensor(xs, dtype=torch.float32)
|
|||
|
|
y = torch.tensor(ys, dtype=torch.float32)
|
|||
|
|
valid_lens = [min(MAX_SEQ_LEN, sum(1 for row in s if any(v != 0 for v in row))) for s in xs]
|
|||
|
|
for i, s in enumerate(xs):
|
|||
|
|
vl = 1
|
|||
|
|
for j in range(1, MAX_SEQ_LEN):
|
|||
|
|
if any(v != 0 for v in s[j]):
|
|||
|
|
vl = j + 1
|
|||
|
|
valid_lens[i] = max(1, vl)
|
|||
|
|
|
|||
|
|
model = _build_torch_model()
|
|||
|
|
opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
|
|||
|
|
loss_fn = nn.BCEWithLogitsLoss()
|
|||
|
|
|
|||
|
|
n = len(xs)
|
|||
|
|
for ep in range(epochs):
|
|||
|
|
perm = torch.randperm(n)
|
|||
|
|
total_loss = 0.0
|
|||
|
|
steps = 0
|
|||
|
|
for start in range(0, n, batch_size):
|
|||
|
|
idx = perm[start : start + batch_size]
|
|||
|
|
bx = X[idx]
|
|||
|
|
by = y[idx]
|
|||
|
|
vl_batch = [valid_lens[i] for i in idx.tolist()]
|
|||
|
|
opt.zero_grad()
|
|||
|
|
logits = model(bx, vl_batch)
|
|||
|
|
loss = loss_fn(logits, by)
|
|||
|
|
loss.backward()
|
|||
|
|
opt.step()
|
|||
|
|
total_loss += loss.item()
|
|||
|
|
steps += 1
|
|||
|
|
acc = 0.0
|
|||
|
|
with torch.no_grad():
|
|||
|
|
pred = (torch.sigmoid(model(X, valid_lens)) > 0.5).float()
|
|||
|
|
acc = (pred == y).float().mean().item()
|
|||
|
|
print(f"epoch {ep + 1}/{epochs} loss={total_loss / max(steps, 1):.4f} acc={acc:.3f}")
|
|||
|
|
|
|||
|
|
numpy_model = export_torch_to_numpy(model)
|
|||
|
|
out_path.parent.mkdir(parents=True, exist_ok=True)
|
|||
|
|
numpy_model.save_json(out_path)
|
|||
|
|
print(f"已导出 {out_path}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
parser = argparse.ArgumentParser()
|
|||
|
|
parser.add_argument("--epochs", type=int, default=20)
|
|||
|
|
parser.add_argument("--from-db", action="store_true")
|
|||
|
|
parser.add_argument("--out", type=str, default="")
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
out = Path(args.out) if args.out else WEIGHTS_PATH
|
|||
|
|
train_and_export(epochs=args.epochs, from_db=args.from_db, out_path=out)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|