📚 PyTorch实战:用LSTM实现文本风格分类(余华vs路遥)

📝 前言

本周我们将之前学习的序列建模基石:RNN-LSTM-GRU理论落地,用PyTorch实现一个基于LSTM的文本风格分类器——通过书名判断作者是余华还是路遥。
这篇博客完全对应之前的知识点:从文本预处理(词表构建、序列填充),到LSTM的门控机制与细胞状态,再到二分类交叉熵损失与模型训练,完整打通“理论→代码”的逻辑链条。


一、数据流向图

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
                    一次前向传播的数据流向图

┌─────────────────────────────────────────────────────────────┐
│ 原始输入 │
│ 例: ["许三观卖血记", "活着"] │
│ 真实标签: [1, 0] (1=余华, 0=路遥) │
└───────────────┬─────────────────────────────────────────────┘


┌─────────────────────────────────────────────────────────────┐
│ 分词 + 查字典 (word_to_idx) │
│ "许三观卖血记" → [15, 27, 3, 48, 22] │
│ "活着" → [36, 5] │
│ 输出: 不等长索引序列列表 │
└───────────────┬─────────────────────────────────────────────┘


┌─────────────────────────────────────────────────────────────┐
│ 填充 (pad_sequence) │
│ [15,27,3,48,22] → [15,27,3,48,22] │
│ [36,5] → [36,5, 0, 0, 0] (0=PAD) │
│ 输出形状: [batch, seq_len] 例: [2, 5] │
└───────────────┬─────────────────────────────────────────────┘


┌─────────────────────────────────────────────────────────────┐
│ Embedding 层 │
│ 输入: [batch, seq_len] 例: [2, 5] 每个位置是整数索引 │
│ 输出: [batch, seq_len, embedding_dim] 例: [2, 5, 16] │
│ 语义: 每个词索引被查表替换成一个16维稠密向量 │
└───────────────┬─────────────────────────────────────────────┘


┌─────────────────────────────────────────────────────────────┐
│ LSTM 层 │
│ 输入: [batch, seq_len, embedding_dim] 例: [2, 5, 16] │
│ │
│ 输出三样东西: │
│ lstm_out: [batch, seq_len, hidden_dim] 例: [2, 5, 32] │
│ 每一步的输出,分类任务不需要它 │
│ │
│ hidden: [num_layers, batch, hidden_dim] 例: [1, 2, 32] │
│ 最后一步的短期记忆,汇总了全句信息,√ 取这个 │
│ │
│ cell: [num_layers, batch, hidden_dim] 例: [1, 2, 32] │
│ 最后一步的长期记忆,分类时不需要 │
└───────────────┬─────────────────────────────────────────────┘


┌─────────────────────────────────────────────────────────────┐
│ 取最后时间步的 hidden,并去掉层维度 (squeeze) │
│ │
│ 选中: hidden (丢弃 lstm_out 和 cell) │
│ 输入: [num_layers, batch, hidden_dim] 例: [1, 2, 32] │
│ 操作: hidden.squeeze(0) 或 hidden[-1] │
│ 输出: [batch, hidden_dim] 例: [2, 32] │
│ │
│ 语义: 把两条句子的"整体理解"分别浓缩成32维向量 │
└───────────────┬─────────────────────────────────────────────┘


┌─────────────────────────────────────────────────────────────┐
│ Linear 层 (全连接) │
│ 输入: [batch, hidden_dim] 例: [2, 32] │
│ 操作: y = Wx + b (W形状 [1, 32], b形状 [1]) │
│ 输出: [batch, 1] 例: [2, 1] │
│ │
│ 语义: 把32维句子理解加权求和成1个原始分数 │
│ 分越高 → 越像余华,分越低 → 越像路遥 │
└───────────────┬─────────────────────────────────────────────┘


┌─────────────────────────────────────────────────────────────┐
│ Sigmoid 层 │
│ 输入: [batch, 1] 例: [[2.3], [-1.1]] │
│ 操作: 1 / (1 + e^(-x)) │
│ 输出: [batch, 1] 例: [[0.91], [0.25]] │
│ │
│ 语义: │
│ > 0.5 → 预测为余华 (1) │
│ ≤ 0.5 → 预测为路遥 (0) │
└───────────────┬─────────────────────────────────────────────┘


┌─────────────────────────────────────────────────────────────┐
│ BCELoss (仅训练时,推理时不走这一步) │
│ 输入: 预测值 [[0.91], [0.25]] , 真实标签 [[1], [0]] │
│ 输出: 损失值 (一个标量) 例: 0.17 │
│ │
│ 语义: 预测越准 → Loss 越小 │
└─────────────────────────────────────────────────────────────┘

二、核心参数解析

参数名出现位置含义示例值作用说明
batch全程批次大小,一次同时处理多少条样本2同时处理两条句子,GPU并行加速,梯度更新更稳定
seq_len填充后序列长度,每条样本被统一到的词数5所有句子补齐到相同长度,短句用PAD填充
vocab_sizeEmbedding词表大小,一共多少个不同的字取决于语料决定Embedding层的查表范围
embedding_dimEmbedding词向量维度,每个字用几个浮点数表示16维度越大表达能力越强,但训练越慢
hidden_dimLSTM, Linear隐藏层维度,LSTM记忆向量的宽度32贯穿LSTM输出和Linear输入,是模型的”脑容量”
num_layersLSTMLSTM堆叠层数1单层适合小数据,多层适合复杂任务
output_dimLinear输出维度1二分类只输出一个分数,多分类则输出类别数
padding_idxEmbedding填充符索引,指定哪个索引不做梯度更新0让PAD符号不参与训练,避免噪声干扰

关键传递关系:

embedding_dim 是 Embedding 的输出维度,也是 LSTM 的输入维度

hidden_dim 是 LSTM 的输出维度,也是 Linear 的输入维度,必须一致

vocab_size 决定 Embedding 的查表范围

num_layers 决定 hidden 和 cell 的第0维大小

三、完整代码及注释

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
import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn.utils.rnn import pad_sequence # 用于把不同长度的句子填充成相同长度

# ============================
# 1. 数据准备
# ============================
# 训练数据:(文本句子, 标签)
# 标签 1 = 余华风格
# 标签 0 = 路遥风格
train_data = [
("许三观卖血记", 1),
("活着", 1),
("我们生活在巨大的差距里", 1),
("平凡的世界", 0),
("人生啊", 0),
("孙少平站在矿井口", 0),
]

# 构建【字符 → 索引】的字典
# 索引 0 固定给 <PAD> 填充符,所有短句子统一长度时用 0 补齐
word_to_idx = {"<PAD>": 0}

# 遍历所有句子,给每个没见过的字分配一个唯一数字ID
for s, _ in train_data:
for w in s: # 按【字】切分,每个字是一个token
if w not in word_to_idx: # 如果这个字不在字典里,就给它分配新索引
word_to_idx[w] = len(word_to_idx) # 新索引 = 当前字典长度(自动递增1、2、3...)

vocab_size = len(word_to_idx) # 词汇表大小(总共有多少个不同的字)

# ============================
# 把句子转换成【数字索引序列】
# sequences 里每一项是一个句子对应的数字列表
sequences = [torch.tensor([word_to_idx[w] for w in s]) for s, _ in train_data]

# 把所有标签单独拿出来,转成浮点型张量(BCELoss 要求)
labels = torch.tensor([l for _, l in train_data], dtype=torch.float32)

# 关键:把不同长度的句子填充成【相同长度】,否则模型无法训练
# batch_first=True → 输出形状 [批次大小, 序列长度]
# padding_value=0 → 用 <PAD> 对应的 0 填充
padded = pad_sequence(sequences, batch_first=True, padding_value=word_to_idx["<PAD>"])
# padded 最终形状: [6, max_seq_len],6个句子,每个句子长度统一为最长句长度

# ============================
# 2. 定义模型:LSTM 文本分类器
# ============================
class WenFengClassifier(nn.Module):
def __init__(self, vocab_size, embedding_dim, hidden_dim, num_layers=1):
super().__init__()

# ======================
# 1. 嵌入层:把【数字索引】转成【低维向量】,让模型能理解语义
# 输入形状:[batch, seq_len]
# 输出形状:[batch, seq_len, embedding_dim]
# padding_idx=0:告诉模型索引0是填充符,不参与学习
# ======================
self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)

# ======================
# 2. LSTM层:处理序列数据,捕捉上下文语义
# 输入维度 = embedding_dim(嵌入层输出)
# 隐藏层维度 = hidden_dim
# batch_first=True:输入形状优先为 [批次, 序列长度, 特征]
# ======================
self.lstm = nn.LSTM(embedding_dim, hidden_dim,
num_layers=num_layers, batch_first=True)

# ======================
# 3. 全连接层:把 LSTM 输出的特征映射成 1 个分数(二分类)
# 输入:hidden_dim 维的句子特征
# 输出:1 个值(未归一化的概率)
# ======================
self.fc = nn.Linear(hidden_dim, 1)

# ======================
# 4. Sigmoid:把分数压缩到 0~1 之间,代表概率
# ======================
self.sigmoid = nn.Sigmoid()

def forward(self, x):
"""
前向传播:数据在模型里的流动过程
x 输入形状:[batch_size, seq_len]
"""
# ① 嵌入层:数字 → 向量
emb = self.embedding(x) # 输出形状:[batch, seq_len, embedding_dim]

# ② LSTM 前向计算
# lstm_out:每个时间步的输出 [batch, seq_len, hidden_dim]
# hidden:最后一步的隐藏状态 [num_layers, batch, hidden_dim]
# cell:最后一步的细胞状态(长时记忆)
lstm_out, (hidden, cell) = self.lstm(emb)

# ③ 取最后一层的 hidden 作为整个句子的特征
# squeeze(0):去掉 num_layers 维度 → [batch, hidden_dim]
final = hidden.squeeze(0)

# ④ 全连接层:特征 → 1个原始分数
logit = self.fc(final) # 输出形状:[batch, 1]

# ⑤ Sigmoid:原始分数 → 0~1 概率
prob = self.sigmoid(logit) # 输出形状:[batch, 1]

return prob # 返回每个句子属于“余华”的概率

# ============================
# 3. 训练配置
# ============================
EMBEDDING_DIM = 16 # 每个字转成16维向量
HIDDEN_DIM = 32 # LSTM 隐藏层大小
EPOCHS = 300 # 训练轮数
LR = 0.01 # 学习率

# 实例化模型
model = WenFengClassifier(vocab_size, EMBEDDING_DIM, HIDDEN_DIM)

# 损失函数:二分类交叉熵(输入必须是 0~1 之间的概率)
criterion = nn.BCELoss()

# 优化器:Adam,自动更新模型参数
optimizer = optim.Adam(model.parameters(), lr=LR)

# ============================
# 4. 开始训练
# ============================
print("==== 开始训练 ====")
for epoch in range(EPOCHS):
model.train() # 切换为训练模式
optimizer.zero_grad() # 梯度清零,避免累积

preds = model(padded).squeeze(1) # 前向传播,[6,1] → 压缩成 [6]
loss = criterion(preds, labels) # 计算损失:预测值 vs 真实标签

loss.backward() # 反向传播,计算梯度
optimizer.step() # 更新参数

# 每30轮打印一次:损失值 & 准确率
if (epoch + 1) % 30 == 0:
acc = (torch.round(preds) == labels).float().mean().item() * 100
print(f"Epoch {epoch+1:3d} | Loss: {loss.item():.4f} | Acc: {acc:.1f}%")

# ============================
# 5. 推理预测(测试模型)
# ============================
def predict(model, sentence):
model.eval() # 切换为评估模式(禁用dropout、batchnorm等)
with torch.no_grad(): # 推理时不计算梯度,节省内存/加速
# 句子 → 索引列表,不在词表的字用 0(<PAD>)代替
idxs = [word_to_idx.get(w, 0) for w in sentence]
tensor = torch.tensor(idxs).unsqueeze(0) # 增加 batch 维度 → [1, seq_len]
prob = model(tensor).item() # 得到预测概率
# 概率>0.5判为余华,否则路遥
return "余华" if prob > 0.5 else "路遥", prob

# 测试两句典型句子
print("\n==== 测试结果 ====")
for s in ["许三观卖血记", "平凡的世界"]:
author, prob = predict(model, s)
print(f"「{s}」→ 预测作者:{author} (概率:{prob:.3f})")

四、模型调优:可调参数与影响分析

4.1 结构参数(影响模型容量)

参数当前值增大影响减小影响调整信号
embedding_dim16表达能力增强,训练变慢,可能过拟合训练加快,表达能力减弱训练集准确率高而测试集低 → 减小;训练集一直上不去 → 增大
hidden_dim32记忆容量增大,可处理更复杂句子训练加快,显存占用减少长句判断差 → 增大;显存不足/训练太慢 → 减小
num_layers1模型变深,能学更抽象特征模型变浅,训练加速数据量小 → 保持1层;数据量大且欠拟合 → 增加到2-3层

4.2 训练参数(影响收敛过程)

参数当前值增大影响减小影响调整信号
lr (学习率)0.01学习加快,但可能不收敛或震荡学习减慢,收敛更稳但训练更久loss忽大忽小 → 减小;loss下降极慢 → 增大;loss变NaN → 立刻减小
epochs300训练更充分训练更短loss仍在下降 → 增大;loss已平稳或反弹 → 减小或加早停
batch_size6(全量)梯度更准,训练稳定梯度噪声大,泛化可能更好显存不足 → 减小;训练不稳定 → 增大

4.3 典型问题诊断表

训练现象可能原因调整方案
训练集准确率100%,测试集却很差数据太少+模型容量过大,过拟合减小embedding_dim/hidden_dim;增加数据量
loss一直不降,准确率在50%附近波动模型容量不足或学习率不合适增大hidden_dim;调整lr
loss变成NaN学习率过大导致梯度爆炸将lr减小一个数量级
训练很快达到高准确率但loss仍高模型开始死记硬背减小epochs;增加dropout(此处未加)

4.4 调优操作流程

先固定epochs=100,用当前参数跑一次,记录loss曲线和最终准确率

观察曲线特征,对照4.3诊断表定位问题类型

每次只调一个参数,改后重跑对比

确认当前参数最优后,再增大epochs做最终训练