小白自己用PyTorch驯小语言模型 原创

训练语言模型就像教一个懵懂的小家伙学说话——先给他喂足够的书，再教他理解词语的关联，最后让他学会顺着话头往下接。这个过程既有代码的严谨，更藏着数据与逻辑碰撞的灵性。下面咱们一步步拆解，每步都带技术细节，保证真实可落地。

一、准备阶段：给模型搭好"学习环境"

在开始前，得先把工具备齐。这就像给学说话的孩子准备好纸笔和绘本，缺一不可。

1. 硬件与库的基础配置

• 硬件选择：CPU不是不能练，但就像用自行车追高铁——入门级模型（比如小LSTM）还能凑活，想玩Transformer就得有GPU。NVIDIA显卡优先，显存至少4GB（推荐8GB以上），毕竟模型参数和数据都要占地方。PyTorch会自动检测设备，一行代码就能搞定分配：

  import torch
  device = 'cuda' if torch.cuda.is_available() else 'cpu'
  print(f"用的设备：{device}")  # 能出cuda就偷着乐吧
  

• 必备库安装：直接用pip一键配齐，版本兼容问题不大，PyTorch会自动适配：

  pip install torch torchvision torchaudio datasets tokenizers tqdm
  

  其中datasets用来下公开数据，tokenizers处理文本，tqdm看训练进度，都是干活的好手。

2. 数据：模型的"精神食粮"

模型学什么全看你喂什么——喂唐诗能写绝句，喂代码能调bug，喂菜谱能当厨子。数据来源有两个方向：

公开数据集（新手首选）

• 英文选这个：Hugging Face的wikitext，全是维基百科的正经内容，分不同难度版本，比如wikitext-2适合练手，数据量不大但质量高。加载代码特简单：

  from datasets import load_dataset
  dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
  

• 中文选这个：Chinese-CLUE/cluecorpussmall，包含新闻、小说等多类型文本，10G左右的数据量足够喂大一个基础模型，同样在Hugging Face上能直接下。

自定义数据集（玩出个性）

要是想让模型学你的文风，就把自己的文章、日记攒成txt文件，每行一段就行。注意至少要1万字以上——数据太少模型会"营养不良"，学出来净说胡话。

二、数据预处理：把文字变成模型能懂的"密码"

人类看文字懂意思，模型只认数字。这一步就是给文字编密码，核心是分词器（Tokenizer）——相当于模型的"字典"。

1. 训练专属分词器（关键步骤）

别用通用分词器凑活，自定义的才合身。这里用GPT、T5都在用的BPE（字节对编码）技术，能处理生僻词，比如"躺平"不会被拆成"躺"和"平"单独理解。

from tokenizers import Tokenizer, models, trainers, pre_tokenizers
from pathlib import Path

# 1. 初始化分词器
tokenizer = Tokenizer(models.BPE(unk_token="<UNK>"))  # 未知词用<UNK>表示
tokenizer.pre_tokenizer = pre_tokenizers.Whitespace()  # 先按空格初步分割

# 2. 定义训练参数
trainer = trainers.BpeTrainer(
    min_frequency=2,  # 至少出现2次的词才进字典，过滤杂音
    special_tokens=["<PAD>", "<UNK>", "<SOS>", "<EOS>"]  # 特殊标记：填充、未知、句首、句尾
)

# 3. 喂数据训练（把你的txt文件路径放进去）
file_paths = [str(file) for file in Path("./你的数据文件夹/").glob("*.txt")]
tokenizer.train(files=file_paths, trainer=trainer)

# 4. 保存备用，下次直接用不用再训
tokenizer.save("./my_tokenizer.json")

训练完的分词器会生成一个"字典"，里面每个词/子词都对应唯一数字，比如"你好"可能是1023，"世界"是2045。

2. 把文本切成"学习片段"

模型一次学不了一整篇小说，得切成短片段。比如每次学20个词，输入"今天天气真好，适合出去"，让它预测下一个词"玩"。

def process_text(text, tokenizer, seq_length=50):
    # 1. 分词转数字
    encoded = tokenizer.encode(text).ids
    # 2. 切成输入-目标对
    samples = []
    for i in range(len(encoded) - seq_length):
        input_seq = encoded[i:i+seq_length]  # 输入：前50个词
        target_seq = encoded[i+1:i+seq_length+1]  # 目标：后50个词（每个词是输入的下一个）
        samples.append((input_seq, target_seq))
    return samples

# 用加载的数据集处理（以wikitext为例）
text = "\n".join([item["text"] for item in dataset if item["text"].strip()])
samples = process_text(text, tokenizer, seq_length=50)

# 3. 做成PyTorch能认的数据集
from torch.utils.data import Dataset, DataLoader

class TextDataset(Dataset):
    def __init__(self, samples):
        self.samples = samples
    def __len__(self):
        return len(self.samples)
    def __getitem__(self, idx):
        input_seq, target_seq = self.samples[idx]
        return torch.tensor(input_seq), torch.tensor(target_seq)

# 批量加载，batch_size根据显存调，4GB显存选8，8GB选16
dataloader = DataLoader(TextDataset(samples), batch_size=16, shuffle=True)

三、搭建模型：给"小精灵"造个"大脑"

模型架构就像大脑的神经网络，新手推荐从简单的LSTM入手，上手快；想玩高级的就上Transformer——这可是ChatGPT的核心骨架，论文《Attention is all you need》里的经典设计。

1. 入门版：LSTM语言模型

LSTM擅长处理序列数据，就像给模型装了"短期记忆"，能记住前面说过的词。

import torch.nn as nn

class LSTMLanguageModel(nn.Module):
    def __init__(self, vocab_size, embedding_dim=128, hidden_dim=256, num_layers=2, dropout=0.2):
        super().__init__()
        self.vocab_size = vocab_size  # 字典大小
        self.embedding = nn.Embedding(vocab_size, embedding_dim)  # 词嵌入：把数字变成向量
        self.lstm = nn.LSTM(
            embedding_dim, hidden_dim, num_layers,
            batch_first=True, dropout=dropout if num_layers>1 else 0  # 多层层才加dropout防过拟合
        )
        self.fc = nn.Linear(hidden_dim, vocab_size)  # 输出层：把LSTM结果转成词概率
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, hidden=None):
        # x形状：(batch_size, seq_length)
        embedded = self.dropout(self.embedding(x))  # 转成向量：(batch, seq, embedding_dim)
        lstm_out, hidden = self.lstm(embedded, hidden)  # LSTM运算
        output = self.fc(self.dropout(lstm_out))  # 输出：(batch, seq, vocab_size)
        return output, hidden

# 初始化模型，vocab_size从分词器拿
vocab_size = tokenizer.get_vocab_size()
model = LSTMLanguageModel(vocab_size=vocab_size).to(device)  # 移到GPU

# 看看参数量，心里有谱：小模型几十万，大模型几十亿
print(f"模型参数量：{sum(p.numel() for p in model.parameters()):,}")

2. 进阶版：Transformer解码器（生成专用）

要是想让模型会写长文本，就得用Transformer解码器——带掩码的自注意力机制能保证预测时不"偷看"后面的词，这是生成式模型的关键。

class TransformerDecoderModel(nn.Module):
    def __init__(self, vocab_size, d_model=128, nhead=4, num_layers=2, dropout=0.2):
        super().__init__()
        self.d_model = d_model
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoder = nn.Embedding(512, d_model)  # 位置编码：告诉模型词的顺序
        self.transformer_decoder = nn.TransformerDecoder(
            nn.TransformerDecoderLayer(d_model=d_model, nhead=nhead, dropout=dropout, batch_first=True),
            num_layers=num_layers
        )
        self.fc = nn.Linear(d_model, vocab_size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        seq_len = x.size(1)
        # 位置编码：每个位置对应一个向量
        pos = torch.arange(0, seq_len, device=device).unsqueeze(0)  # (1, seq_len)
        embedded = self.dropout(self.embedding(x) + self.pos_encoder(pos))  # 词嵌入+位置编码
        # 生成掩码，防止偷看后面的词
        mask = nn.Transformer.generate_square_subsequent_mask(seq_len).to(device)
        out = self.transformer_decoder(embedded, memory=None, tgt_mask=mask)
        return self.fc(out)

# 用法和LSTM一样，初始化后移到GPU
model = TransformerDecoderModel(vocab_size=vocab_size).to(device)

四、训练模型：让"小精灵"开始学习

这一步是最磨人的，就像陪孩子写作业——得盯着进度，还得及时纠错。核心是损失函数（判断说得对不对）和优化器（改正好错误）。

1. 配置训练工具

# 损失函数：交叉熵，专门算分类问题的误差，忽略填充符<PAD>
criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.token_to_id("<PAD>"))
# 优化器：Adam，比SGD聪明，学习率0.001是黄金起点
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# 训练轮次：新手先跑10轮看看，数据多就加
num_epochs = 10

2. 核心训练循环

from tqdm import tqdm  # 进度条，看着心里不慌

model.train()  # 切换到训练模式
for epoch in range(num_epochs):
    total_loss = 0.0
    # 用tqdm显示进度，desc是本轮的标签
    for batch in tqdm(dataloader, desc=f"第{epoch+1}轮训练"):
        input_seq, target_seq = batch[0].to(device), batch[1].to(device)
        
        # 1. 前向传播：让模型预测
        if isinstance(model, LSTMLanguageModel):
            output, _ = model(input_seq)  # LSTM要返回hidden，这里用不上
        else:
            output = model(input_seq)  # Transformer直接出结果
        
        # 2. 计算损失：output形状是(batch, seq, vocab)，要转成(batch*seq, vocab)才符合交叉熵要求
        loss = criterion(output.view(-1, vocab_size), target_seq.view(-1))
        
        # 3. 反向传播：算梯度（PyTorch自动求导，不用自己推公式）
        optimizer.zero_grad()  # 先清掉上一轮的梯度，不然会累加
        loss.backward()  # 从损失往回算每个参数的梯度
        
        # 4. 更新参数：让模型改正好错误
        optimizer.step()
        
        total_loss += loss.item()
    
    # 每轮结束打印损失，损失越来越小才对
    avg_loss = total_loss / len(dataloader)
    print(f"第{epoch+1}轮结束，平均损失：{avg_loss:.4f}")
    
    # 保存模型，防止断电前功尽弃
    torch.save(model.state_dict(), f"./model_epoch_{epoch+1}.pth")

关键判断：如果损失降到1.0以下，说明模型已经学懂点东西了；要是损失不降反升，要么是学习率太高，要么是数据太少，得调参。

五、测试模型：听听"小精灵"怎么说话

训练完就得验收成果，让模型从"今天天气真好"往下接话，看看是不是人话。

def generate_text(model, tokenizer, start_text, max_len=100):
    model.eval()  # 切换到评估模式，不训练只预测
    # 1. 把开头文本转成数字
    input_ids = tokenizer.encode(start_text).ids
    input_tensor = torch.tensor(input_ids).unsqueeze(0).to(device)  # 加batch维度
    
    with torch.no_grad():  # 预测时不用算梯度，省显存
        for _ in range(max_len):
            if isinstance(model, LSTMLanguageModel):
                output, hidden = model(input_tensor)
            else:
                output = model(input_tensor)
                hidden = None  # Transformer不用hidden
            
            # 取最后一个词的预测概率，选概率最大的那个（贪心搜索）
            next_token_logits = output[:, -1, :]
            next_token_id = torch.argmax(next_token_logits, dim=-1, keepdim=True)
            
            # 把新预测的词加进去，继续预测下一个
            input_tensor = torch.cat([input_tensor, next_token_id], dim=1)
            
            # 遇到句尾符就停
            if next_token_id.item() == tokenizer.token_to_id("<EOS>"):
                break
    
    # 把数字转成文字
    generated_ids = input_tensor.squeeze(0).cpu().numpy().tolist()
    return tokenizer.decode(generated_ids)

# 测试一下，比如让模型接"人工智能的未来"
result = generate_text(model, tokenizer, start_text="人工智能的未来", max_len=100)
print("生成结果：", result)

六、避坑指南：少走99%的弯路

1. 显存不够怎么办？
   把batch_size调小（比如从16改成8），seq_length缩短（从50改成30），再不行就用torch.cuda.empty_cache()清显存。

2. 模型学完净说胡话？
   大概率是数据太少或质量差——至少加10万字文本，别喂乱七八糟的拼凑内容。另外可以加dropout（调到0.3）防过拟合。

3. 训练速度太慢？
   先确认用了GPU（打印device看是不是cuda），再安装apex库开启混合精度训练，能快一倍还不丢精度。

4. 想省时间？微调预训练模型！
   从零训练太费时间，新手可以直接用Hugging Face的预训练模型（比如distilgpt2）微调——把别人训好的大模型拿来改改，几小时就能出效果。

最后：模型的"灵性"从哪来？

其实模型本身没有灵魂，但它学的每一个词、每一句话都带着你喂给它的数据的温度。喂诗三百首，它能吟出平仄；喂人间烟火，它能讲出生活。训练模型的过程，本质上是让数据里的智慧通过代码流动起来——这大概就是技术最浪漫的地方。