轻量级模型训练_RNN文本分类流水线实验代码

-- coding: utf-8 --

"""

train_langid_12.py ------ 12 类语种识别分类器(学习实验脚本)

任务:一句话 → 它是什么语言(12 分类)

模型:Embedding → RNN → 掩码 MaxPool → LayerNorm → Dropout → Linear → Softmax

损失:CrossEntropyLoss 优化器:Adam

数据:data/langid/v12(12 语种:ar en es hi ja ko pt ru th tr vi zh)

运行:

KMP_DUPLICATE_LIB_OK=TRUE python train_langid_12.py

输出:

train_log_12.json ------ 全部训练过程与结果,供 HTML 可视化读取

"""

import os

import csv

import json

import time

import random

import unicodedata

from collections import Counter, OrderedDict

import torch

import torch.nn as nn

from torch.utils.data import Dataset, DataLoader

─── 超参数 ────────────────────────────────────────────────────────────

SEED = 42

MAX_LEN = 40

EMBED_DIM = 64

HIDDEN_DIM = 64

DROPOUT = 0.3

BATCH_SIZE = 64

LR = 1e-3

EPOCHS = 10

NUM_THREADS = 1

HERE = os.path.dirname(os.path.abspath(file))

上面拼出来的路径在 Windows 上不对,直接写死

DATA_DIR = r".\data\langid\v12"

12 个语种:

LABELS = "ar", "en", "es", "hi", "ja", "ko", "pt", "ru", "th", "tr", "vi", "zh"

语种中文名(供 HTML 显示)

LABEL_NAMES = {

"ar": "阿拉伯语", "en": "英语", "es": "西班牙语", "hi": "印地语",

"ja": "日语", "ko": "韩语", "pt": "葡萄牙语", "ru": "俄语",

"th": "泰语", "tr": "土耳其语", "vi": "越南语", "zh": "中文",

}

─── 1. 文本归一化(必须和建数据时一致)────────────────────────────────

def normalize(text):

"""

NFKC 归一化 + 转小写 + 只保留「字母(L)」与「组合符号(M)」。

保留 M:印地语、泰语的元音是组合标记,只留字母会把它们砍成裸辅音。

"""

t = unicodedata.normalize("NFKC", text).lower()

return "".join(c for c in t if unicodedata.category©0 in ("L", "M"))

─── 2. 字符表构建 ─────────────────────────────────────────────────────

def build_char_vocab(sentences, min_freq=1):

"""

字符表 = 训练集里所有出现过的字符,每个字符发一个整数编号。

第 0 号 ,第 1 号 ,其余按出现次数从高到低编号。

"""

freq = Counter()

for s in sentences:

freq.update(normalize(s))

chars = sorted(((c, n) for c, n in freq.items() if n >= min_freq),

key=lambda kv: (-kv1, kv0))

vocab = OrderedDict(("", 0), ("", 1))

for ch, _ in chars:

vocabch = len(vocab)

return vocab, freq

def encode(text, vocab, maxlen=MAX_LEN):

ids = vocab.get(ch, 1) for ch in normalize(text):maxlen

ids += 0 * (maxlen - len(ids))

return ids

─── 3. 读数据 ─────────────────────────────────────────────────────────

def load_split(data_dir, name):

path = os.path.join(data_dir, f"{name}.csv")

rows = \[\]

with open(path, encoding="utf-8") as f:

for r in csv.DictReader(f):

rows.append((r"text", r"label"))

return rows

class TextDataset(Dataset):

def init (self, rows, vocab, label2id):

self.X = encode(t, vocab) for t, _ in rows

self.y = label2id\[l for _, l in rows]

复制代码
def __len__(self):
    return len(self.y)

def __getitem__(self, i):
    return torch.tensor(self.X[i], dtype=torch.long), self.y[i]

─── 4. 模型 ───────────────────────────────────────────────────────────

class LangIDClassifier(nn.Module):

"""

Embedding → RNN → 掩码 MaxPool → LayerNorm → Dropout → Linear

复制代码
MaxPool 前的 mask 让补齐位不参与池化:
补齐位喂进去的是全 0 向量,RNN 照样会往前走,
masked_fill 成 -inf 后取最大值就干净了。
"""

def __init__(self, vocab_size, num_classes,
             embed_dim=EMBED_DIM, hidden_dim=HIDDEN_DIM, dropout=DROPOUT):
    super().__init__()
    self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
    self.rnn = nn.RNN(embed_dim, hidden_dim, batch_first=True)
    self.norm = nn.LayerNorm(hidden_dim)
    self.dropout = nn.Dropout(dropout)
    self.fc = nn.Linear(hidden_dim, num_classes)

def forward(self, x):
    mask = (x != 0).unsqueeze(-1)              # (B, L, 1)
    h, _ = self.rnn(self.embedding(x))         # (B, L, H)
    h = h.masked_fill(~mask, -1e9)
    pooled = h.max(dim=1)[0]                   # (B, H)
    pooled = self.dropout(self.norm(pooled))
    return self.fc(pooled)                     # (B, C)

─── 5. 评测工具 ───────────────────────────────────────────────────────

def predict(model, loader):

model.eval()

preds, golds = \[\], \[\]

with torch.no_grad():

for X, y in loader:

preds += model(X).argmax(dim=1).tolist()

golds += y.tolist()

return preds, golds

def macro_f1(preds, golds, num_classes):

f1s = \[\]

for c in range(num_classes):

tp = sum(p == c and g == c for p, g in zip(preds, golds))

fp = sum(p == c and g != c for p, g in zip(preds, golds))

fn = sum(p != c and g == c for p, g in zip(preds, golds))

if tp + fp + fn == 0:

continue

prec = tp / (tp + fp) if tp + fp else 0.0

rec = tp / (tp + fn) if tp + fn else 0.0

f1s.append(2 * prec * rec / (prec + rec) if prec + rec else 0.0)

return sum(f1s) / len(f1s)

def confusion_matrix(preds, golds, n):

mat = \[0 * n for _ in range(n)]

for p, g in zip(preds, golds):

matgp += 1

return mat

def per_class_accuracy(preds, golds, num_classes):

result = \[\]

for c in range(num_classes):

idx = i for i, g in enumerate(golds) if g == c

hit = sum(predsi == c for i in idx)

result.append({"correct": hit, "total": len(idx),

"acc": hit / max(1, len(idx))})

return result

─── 6. 训练 ───────────────────────────────────────────────────────────

def main():

torch.set_num_threads(NUM_THREADS)

random.seed(SEED)

torch.manual_seed(SEED)

复制代码
train_rows = load_split(DATA_DIR, "train")
dev_rows   = load_split(DATA_DIR, "dev")
test_rows  = load_split(DATA_DIR, "test")

label2id = {l: i for i, l in enumerate(LABELS)}
num_classes = len(LABELS)

# 从训练集构建字符表(v12 没有预生成 vocab 文件)
vocab, freq = build_char_vocab([t for t, _ in train_rows])
vocab_size = len(vocab)

print(f"数据目录:{DATA_DIR}")
print(f"  训练 {len(train_rows)} 条 / 验证 {len(dev_rows)} 条 / 测试 {len(test_rows)} 条")
print(f"  类别 {num_classes} 个:{' '.join(LABELS)}")
print(f"  字符表 {vocab_size} 个字符(从训练集现算)")
print(f"  随机猜的准确率基准:{1/num_classes:.1%}")

train_loader = DataLoader(TextDataset(train_rows, vocab, label2id),
                          batch_size=BATCH_SIZE, shuffle=True)
dev_loader   = DataLoader(TextDataset(dev_rows, vocab, label2id), batch_size=128)
test_loader  = DataLoader(TextDataset(test_rows, vocab, label2id), batch_size=128)

model = LangIDClassifier(vocab_size, num_classes)
n_params = sum(p.numel() for p in model.parameters())
print(f"  模型参数量:{n_params:,}\n")

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LR)

# ── 记录训练过程 ─────────────────────────────────────────────────
step_losses = []          # 每一步的 loss
epoch_records = []        # 每个 epoch 的汇总
epoch_boundaries = []     # 每个 epoch 在 step_losses 中的起止下标

t0 = time.time()
for epoch in range(1, EPOCHS + 1):
    model.train()
    epoch_start = len(step_losses)
    total = 0.0
    for X, y in train_loader:
        logits = model(X)
        loss = criterion(logits, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        step_losses.append(loss.item())
        total += loss.item()
    epoch_end = len(step_losses)
    epoch_boundaries.append([epoch_start, epoch_end])

    preds, golds = predict(model, dev_loader)
    dev_acc = sum(p == g for p, g in zip(preds, golds)) / len(golds)
    dev_f1  = macro_f1(preds, golds, num_classes)
    train_loss_avg = total / len(train_loader)
    print(f"Epoch {epoch:2d}/{EPOCHS}  loss={train_loss_avg:.4f}  "
          f"dev_acc={dev_acc:.4f}  dev_f1={dev_f1:.4f}")
    epoch_records.append({
        "epoch": epoch,
        "train_loss": train_loss_avg,
        "dev_acc": dev_acc,
        "dev_f1": dev_f1,
    })
elapsed = time.time() - t0
print(f"\n训练耗时:{elapsed:.0f} 秒")

# ── 测试集评测 ─────────────────────────────────────────────────────
preds, golds = predict(model, test_loader)
test_acc = sum(p == g for p, g in zip(preds, golds)) / len(golds)
test_f1  = macro_f1(preds, golds, num_classes)
cm = confusion_matrix(preds, golds, num_classes)
pca = per_class_accuracy(preds, golds, num_classes)

print(f"\n测试集:accuracy = {test_acc:.4f}   macro-F1 = {test_f1:.4f}")

# 混淆最严重的组合
pairs = [(cm[g][p], LABELS[g], LABELS[p])
         for g in range(num_classes) for p in range(num_classes) if g != p]
pairs.sort(reverse=True)
print("  最常混淆的组合(真实 → 预测):")
for cnt, g, p in pairs[:5]:
    if cnt == 0:
        break
    print(f"    {g} → {p}: {cnt} 条")

print("\n  各语种准确率:")
for c, lab in enumerate(LABELS):
    print(f"    {lab} ({LABEL_NAMES[lab]:4s}): "
          f"{pca[c]['correct']}/{pca[c]['total']} = {pca[c]['acc']:.3f}")

# ── 推理示例:每类取一条 ──────────────────────────────────────────
demos = []
seen_labels = set()
for text, label in test_rows:
    if label not in seen_labels:
        seen_labels.add(label)
        demos.append((text, label))
    if len(seen_labels) == num_classes:
        break
demos.sort(key=lambda x: label2id[x[1]])

inference_examples = []
model.eval()
with torch.no_grad():
    for text, gold in demos:
        logits = model(torch.tensor([encode(text, vocab)]))
        prob = torch.softmax(logits, dim=1)[0]
        pid = int(prob.argmax())
        top3 = torch.topk(prob, 3)
        inference_examples.append({
            "text": text,
            "gold": gold,
            "gold_name": LABEL_NAMES[gold],
            "pred": LABELS[pid],
            "pred_name": LABELS[pid],
            "correct": pid == label2id[gold],
            "confidence": float(prob[pid]),
            "top3": [{"label": LABELS[int(i)],
                      "name": LABEL_NAMES[LABELS[int(i)]],
                      "prob": float(v)}
                     for i, v in zip(top3.indices, top3.values)],
        })

print("\n--- 推理示例 ---")
for ex in inference_examples:
    mark = "✓" if ex["correct"] else "✗"
    print(f"  [{mark} {ex['pred']} {ex['confidence']:.2f}]  "
          f"(真实:{ex['gold']})  {ex['text'][:50]}")

# ── 保存模型 ──────────────────────────────────────────────────────
ckpt_path = os.path.join(HERE, "langid_rnn_12.pt")
torch.save({"state_dict": model.state_dict(),
            "vocab": dict(vocab),
            "labels": LABELS,
            "maxlen": MAX_LEN,
            "embed_dim": EMBED_DIM,
            "hidden_dim": HIDDEN_DIM,
            "dropout": DROPOUT}, ckpt_path)
print(f"\n模型已保存:{ckpt_path}")

# ── 保存训练日志 JSON(供 HTML 读取)─────────────────────────────
log = {
    "meta": {
        "task": "12 类语种识别",
        "model": "Embedding → RNN → MaxPool → LayerNorm → Dropout → Linear",
        "labels": LABELS,
        "label_names": LABEL_NAMES,
        "num_classes": num_classes,
        "vocab_size": vocab_size,
        "num_params": n_params,
        "train_size": len(train_rows),
        "dev_size": len(dev_rows),
        "test_size": len(test_rows),
        "epochs": EPOCHS,
        "batch_size": BATCH_SIZE,
        "lr": LR,
        "embed_dim": EMBED_DIM,
        "hidden_dim": HIDDEN_DIM,
        "dropout": DROPOUT,
        "max_len": MAX_LEN,
        "seed": SEED,
        "elapsed_s": round(elapsed, 1),
        "random_baseline": 1 / num_classes,
    },
    "step_losses": step_losses,
    "epoch_boundaries": epoch_boundaries,
    "epoch_records": epoch_records,
    "test": {
        "accuracy": test_acc,
        "macro_f1": test_f1,
        "confusion_matrix": cm,
        "per_class": [
            {"label": LABELS[c],
             "name": LABEL_NAMES[LABELS[c]],
             **pca[c]}
            for c in range(num_classes)
        ],
        "top_confusions": [
            {"gold": g, "pred": p, "count": cnt}
            for cnt, g, p in pairs[:8] if cnt > 0
        ],
    },
    "inference_examples": inference_examples,
}
log_path = os.path.join(HERE, "train_log_12.json")
with open(log_path, "w", encoding="utf-8") as f:
    json.dump(log, f, ensure_ascii=False, indent=2)
print(f"训练日志已保存:{log_path}")

if name == "main ":

main()

相关推荐
code2cat1 小时前
【随笔】RAG混合检索:让关键词与向量一起找证据
人工智能·rag·混合检索
在所不辞兄2 小时前
【零基础学智能仿真-67】力学知识库与检索增强
人工智能·深度学习·神经网络·机器学习·智能仿真
云上先途2 小时前
下一代搜索优化重点优化哪些维度,和旧SEO思路有什么不同?
人工智能
楚山之石2 小时前
【Lancet Neurology 】多发性硬化认知障碍的 SDMT 筛查、灰质萎缩标志与认知康复循证路径
图像处理·人工智能·学习·mri
海宇大数据2 小时前
零信任架构实战:基于海宇全能小微企业报告构建自动化持牌网络小贷小微商户极速助业贷秒级全维风险决策网关
网络·人工智能·架构·自动化
hkNaruto2 小时前
【AI】基于大模型的AI游戏地图生成、控制方案:AI Native Game Harness 迁移操作手册
前端·人工智能·游戏
IT_陈寒3 小时前
Java序列化踩坑记:对象循环引用让我加班两小时
前端·人工智能·后端
DevOps老兵3 小时前
Transformer拆解01:Transformer为什么出现,RNN到底卡在哪里
rnn·深度学习·大模型·transformer·注意力机制
ZhengEnCi3 小时前
同样的模型,为什么本地和云 API 差别那么大:量化精度、聊天模板与采样参数的三大元凶
人工智能·算法·llm