30,PyTorch 序列模型的构建与训练-CSDN博客

在这里插入图片描述

30, PyTorch 序列模型的构建与训练

在上一节中，我们已经用 Hugging Face tokenizers 训练出了一份垂直领域专用的 BPE 分词器，并把任意文本压缩成了短、准、省的张量序列。现在，是时候把这些张量喂给真正的序列模型，完成「从字符到语义」的最后一跃。

本节聚焦「如何用最精简的代码在 PyTorch 里搭建并训练一个可落地的序列模型」。无论你是想跑通一个 LSTM 基线，还是想实现一个 1-D GAN 做数据增强，抑或想微调一个 Transformer Encoder 做下游分类，都可以直接套用本节模板。

30.1 统一流水线回顾

整体流程与 28、29 节保持一致：

原始文本
   ↓ 29 节自定义 Tokenizer
离散 token id
   ↓ 本节 Dataset / DataLoader
张量批次
   ↓ 本节 Model
logits / loss
   ↓ backward
更新梯度

因此，本节 Dataset 与 29.4 完全兼容，无需修改数据侧代码。

30.2 环境

pip install -U torch==2.2 datasets==2.18 accelerate==0.30

无需 transformers，本节全部用裸 PyTorch 实现，方便你魔改。

30.3 Dataset & DataLoader

沿用 29 节的 corpus.txt，每行一条样本。我们做一个「下一 token 预测」的自回归任务。

from datasets import load_dataset
from torch.utils.data import DataLoader
from functools import partial

tokenizer = PreTrainedTokenizerFast(tokenizer_file="my_bpe.json")

def tokenize(examples):
    out = tokenizer(
        examples["text"],
        truncation=True,
        max_length=128,
    )
    # labels 右移一位做自回归
    out["labels"] = out["input_ids"][1:] + [tokenizer.eos_token_id]
    return out

ds = load_dataset("text", data_files={"train": "data/corpus.txt"})["train"]
ds = ds.map(tokenize, batched=True, remove_columns=["text"])
ds.set_format(type="torch", columns=["input_ids", "labels"])

loader = DataLoader(ds, batch_size=64, shuffle=True, num_workers=4)

30.4 模型设计

30.4.1 LSTM 基线（5 行）

import torch.nn as nn

class LSTMModel(nn.Module):
    def __init__(self, vocab_size, d_model=512, n_layers=3):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.lstm  = nn.LSTM(d_model, d_model, n_layers, batch_first=True)
        self.head  = nn.Linear(d_model, vocab_size)

    def forward(self, x, labels=None):
        x = self.embed(x)
        logits, _ = self.lstm(x)
        logits = self.head(logits)
        if labels is not None:
            loss = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)(
                logits.view(-1, logits.size(-1)), labels.view(-1))
            return logits, loss
        return logits

30.4.2 Transformer Encoder（轻量版）

如果你更关心并行训练速度，可以直接上 Transformer：

from torch.nn import TransformerEncoder, TransformerEncoderLayer

class TinyTransformer(nn.Module):
    def __init__(self, vocab_size, d_model=512, nhead=8, nlayers=6):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.pos   = nn.Parameter(torch.randn(1, 128, d_model))
        encoder_layer = TransformerEncoderLayer(d_model, nhead, 4*d_model, batch_first=True)
        self.encoder = TransformerEncoder(encoder_layer, nlayers)
        self.head = nn.Linear(d_model, vocab_size)

    def forward(self, x, labels=None):
        x = self.embed(x) + self.pos[:, :x.size(1)]
        mask = nn.Transformer.generate_square_subsequent_mask(x.size(1)).to(x.device)
        logits = self.head(self.encoder(x, mask))
        if labels is not None:
            loss = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)(
                logits.view(-1, logits.size(-1)), labels.view(-1))
            return logits, loss
        return logits

30.5 训练循环（单卡 / DDP 通杀）

import torch, time
from accelerate import Accelerator

accelerator = Accelerator()
model = LSTMModel(len(tokenizer)).to(accelerator.device)
optimizer = torch.optim.AdamW(model.parameters(), 1e-3)

model, optimizer, loader = accelerator.prepare(model, optimizer, loader)

for epoch in range(3):
    t0 = time.time()
    for step, batch in enumerate(loader):
        logits, loss = model(**batch)
        accelerator.backward(loss)
        optimizer.step(); optimizer.zero_grad()
    print(f"epoch {epoch} | loss {loss.item():.4f} | t {(time.time()-t0):.1f}s")

单卡：直接 python train.py

多卡：

accelerate config  # 首次交互式配置
accelerate launch train.py

30.6 断点续训 & 模型保存

accelerator.save_state("ckpt")
# 恢复
accelerator.load_state("ckpt")

导出最终权重：

torch.save(model.state_dict(), "lstm_final.pt")

30.7 推理示例

model.eval()
prompt = "血常规白细胞计数"
ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
    for _ in range(20):
        logits = model(ids)
        next_id = logits[0, -1].argmax().unsqueeze(0)
        ids = torch.cat([ids, next_id.unsqueeze(0)], dim=1)
print(tokenizer.decode(ids[0]))

30.8 性能基准（RTX 4090）

模型	参数量	速度 (tok/s)	显存 (batch=64)
LSTM-3	47 M	120 k	4.1 GB
TinyTransformer-6	86 M	185 k	5.9 GB

结论：Transformer 并行优势明显，但 LSTM 在小词表、低显存场景依旧能打。

30.9 一键脚本

项目已提供 train_lm.py：

python train_lm.py \
    --model_type lstm \
    --tokenizer_path my_bpe.json \
    --data_path data/corpus.txt \
    --max_len 128 \
    --batch_size 64 \
    --lr 1e-3 \
    --epochs 3 \
    --save_dir ./ckpt

支持 --model_type transformer，自动切换。