从零复现LSTM天池新闻文本分类：一份能跑通的Python源码

简介：这份Python源码包围绕天池新闻文本分类比赛展开，采用LSTM作为核心模型，适合人工智能、计算机及相关专业学生、教师与企业员工用于课程设计、毕业设计或赛题复现。包内共25个文件，以14个py脚本为主体，辅以9个pyc编译文件、1个txt词表与1个json配置，压缩包约58KB，结构紧凑，便于快速阅读与二次开发。代码涵盖LSTM编码器、TextCNN编码器、注意力机制、BERT编码器、数据预处理、模型训练与对抗训练等模块，并配有训练入口脚本，可帮助读者理解从文本向量化到模型评估的完整流程。已有161人学习，说明该方案具备一定参考价值。对于希望掌握新闻文本分类实战、对比不同编码器效果或在此基础上迁移到其他NLP任务的读者，这份源码能提供可运行的基线实现与清晰的模块划分，降低从零搭建的门槛。

1. 从零复现 LSTM 天池新闻文本分类：一份能跑通的 Python 源码该长什么样

天池新闻文本分类这个赛题，很多人第一次跑 LSTM 都会卡在同一个地方：代码能跑，但分数上不去，或者干脆在词表构建那一步就报内存错误。我见过太多人把python环境配好、pycharm配置完，拿到一份源码直接python train.py，结果要么是KeyError，要么是 loss 不降，最后不了了之。这份基于 LSTM 的新闻文本分类源码，核心要解决的就是把 14 类新闻标题和正文映射成固定长度序列，用 Embedding + LSTM + 全连接做多分类。它适合已经会python基础语法、想拿一个完整 NLP 项目练手的人，也适合想搞懂「为什么我的 LSTM 比别人的低 5 个点」的熟手。下面我按自己复现时的顺序，把数据、模型、训练、调参和踩坑一条条拆开。

2. 数据读取与词表构建：别让内存和 OOV 拖垮你的 LSTM

2.1 天池新闻数据的真实结构和读取方式

天池新闻文本分类的原始数据通常是train_set.csv和test_a.csv，每行包含label和text两列，text是新闻标题加正文的拼接，用空格分词后的形式。很多人直接用pd.read_csv读，然后text.split()做词表，这在数据量小的时候没问题，但天池这个赛题训练集有 20 万条，每条平均几百个词，全量加载后内存占用很容易超过 8G。我一般会先看数据分布，再决定要不要做截断。

import pandas as pd from collections import Counter # 读取训练集，指定分隔符和列名 train_df = pd.read_csv('train_set.csv', sep='\t') # 查看类别分布，确认是否均衡 print(train_df['label'].value_counts().sort_index()) # 统计每条文本的词数分布，决定 max_len text_len = train_df['text'].apply(lambda x: len(x.split())) print(text_len.describe(percentiles=[0.5, 0.9, 0.95, 0.99]))

这段代码先确认标签是否从 0 到 13 连续，再通过分位数看 95% 的文本长度落在哪里。如果 95% 分位数是 800，那max_len设 1000 就够，设 2000 只会让 LSTM 的序列过长，梯度回传变慢，显存也吃紧。参数上，sep='\t'是天池数据常见的制表符分隔，如果实际是逗号，改成sep=','即可。percentiles列表里我习惯加 0.99，防止极端长文本影响判断。

2.2 词表构建的两种策略和 OOV 处理

词表构建直接决定 Embedding 层的输入质量。常见做法是统计所有训练文本的词频，取 top N 个词，剩下的映射为<UNK>。但这里有个坑：如果只对训练集建词表，测试集里出现的新词全变<UNK>，模型在验证集上会掉点。我一般会把训练集和测试集合并后一起统计词频，再切分。另外，词表大小不是越大越好，天池这个赛题词表控制在 5 万到 10 万之间比较稳，太大 Embedding 参数量暴涨，太小 OOV 太多。

from collections import Counter # 合并训练和测试的文本，统一建词表 test_df = pd.read_csv('test_a.csv', sep='\t') all_text = pd.concat([train_df['text'], test_df['text']], ignore_index=True) # 统计词频，只保留出现次数 >= 2 的词 word_counter = Counter() for text in all_text: word_counter.update(text.split()) # 构建词表：0 留给 padding，1 留给 UNK vocab = {'<PAD>': 0, '<UNK>': 1} for word, count in word_counter.most_common(): if count < 2: break if len(vocab) >= 100000: break vocab[word] = len(vocab) print(f'词表大小: {len(vocab)}')

这里count < 2过滤掉只出现一次的词，能显著减小词表且对精度影响很小。len(vocab) >= 100000是硬上限，防止个别高频噪声词撑爆词表。建完词表后，把文本转成 id 序列时，遇到不在词表里的词就填 1。注意<PAD>和<UNK>的 id 必须固定，后面 Embedding 层要对应。

2.3 把文本转成定长 id 序列的完整函数

有了词表，下一步是把每条文本转成固定长度的 id 列表。短了补 0，长了截断。截断策略有从头部截、从尾部截、头尾各截一半，新闻文本的关键信息通常在开头，我一般保留前max_len个词。

def text_to_ids(text, vocab, max_len=1000): """将文本转为定长 id 序列，不足补 0，超出截断""" words = text.split() ids = [vocab.get(w, 1) for w in words] # 1 是 UNK if len(ids) < max_len: ids = ids + [0] * (max_len - len(ids)) else: ids = ids[:max_len] return ids # 应用到训练集和测试集 train_df['ids'] = train_df['text'].apply(lambda x: text_to_ids(x, vocab)) test_df['ids'] = test_df['text'].apply(lambda x: text_to_ids(x, vocab))

vocab.get(w, 1)保证 OOV 词映射到<UNK>。补 0 操作放在后面，因为 LSTM 对 padding 位置可以通过pack_padded_sequence忽略，但为了简单，很多源码直接补 0 后接Embedding再送 LSTM，效果也能接受。如果显存够，max_len可以设 1200 左右，再大收益递减。

3. LSTM 模型搭建：从 Embedding 到分类头的每一层参数

3.1 模型整体结构和前向传播逻辑

这份源码的模型部分通常是一个nn.Module，包含 Embedding 层、LSTM 层、全连接层。Embedding 把 id 序列映射成稠密向量，LSTM 提取序列特征，取最后一个时间步的输出或做平均池化，再经过全连接映射到 14 类。我一般会在 LSTM 后加一个 Dropout，防止过拟合。下面是一个可直接用的模型定义。

import torch import torch.nn as nn class LSTMClassifier(nn.Module): def __init__(self, vocab_size, embed_dim=300, hidden_dim=256, num_layers=2, num_classes=14, dropout=0.3): super().__init__() # padding_idx=0 表示 0 对应的 embedding 不参与梯度更新 self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0) self.lstm = nn.LSTM(embed_dim, hidden_dim, num_layers, batch_first=True, bidirectional=True, dropout=dropout if num_layers > 1 else 0) self.dropout = nn.Dropout(dropout) self.fc = nn.Linear(hidden_dim * 2, num_classes) # 双向所以乘 2 def forward(self, x): # x: (batch, seq_len) emb = self.embedding(x) # (batch, seq_len, embed_dim) out, (h, c) = self.lstm(emb) # out: (batch, seq_len, hidden*2) # 取最后一个时间步的输出 out = out[:, -1, :] out = self.dropout(out) logits = self.fc(out) return logits

padding_idx=0让 padding 的 embedding 保持为 0 且不更新，减少噪声。bidirectional=True让 LSTM 同时看前后文，对新闻分类这种任务通常比单向高 1 到 2 个点。num_layers=2时 LSTM 内部会加 dropout，但只有层数大于 1 才生效。hidden_dim * 2是因为双向拼接。取out[:, -1, :]是取最后一个时间步，如果序列补了很多 0，最后一个时间步可能是 padding，这时改用平均池化更稳。

3.2 关键参数怎么设：embed_dim、hidden_dim、num_layers

这三个参数直接决定模型容量和训练速度。embed_dim常见 128、256、300，天池这个赛题用 300 预训练词向量初始化效果更好，但源码里如果没提供预训练向量，从随机初始化开始，300 和 256 差别不大。hidden_dim我一般设 256，双向后输出 512，再大显存吃紧且容易过拟合。num_layers设 2 足够，3 层以上训练慢且提升有限。

如果训练集 loss 降得很快但验证集不降，先把dropout调到 0.5，再把hidden_dim降到 128。如果训练集 loss 都降不下去，检查词表是不是太小，或者max_len截断太狠。

3.3 用 pack_padded_sequence 处理变长序列的正确姿势

上面模型对 padding 位置也做了 LSTM 计算，虽然padding_idx=0减少了 embedding 噪声，但 LSTM 仍会在 padding 上消耗计算。更规范的做法是用pack_padded_sequence，只对真实长度做 LSTM。这需要先按长度排序，再 pack，最后 unpack。

from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence def forward_with_pack(self, x, lengths): emb = self.embedding(x) # 按长度降序排序，pack 要求 lengths, sort_idx = lengths.sort(descending=True) emb = emb[sort_idx] packed = pack_padded_sequence(emb, lengths.cpu(), batch_first=True) packed_out, (h, c) = self.lstm(packed) out, _ = pad_packed_sequence(packed_out, batch_first=True) # 还原原始顺序 _, unsort_idx = sort_idx.sort() out = out[unsort_idx] out = out[:, -1, :] return self.fc(self.dropout(out))

lengths必须是 CPU 上的 int64 张量。排序后要记住sort_idx，最后用unsort_idx还原，否则 batch 内顺序错乱，loss 会异常。如果嫌麻烦，直接补 0 送 LSTM 也能跑，但显存占用会高 20% 左右。

4. 训练循环与验证：让 loss 真正降下来的几个开关

4.1 数据加载和 batch 划分

训练循环第一步是把 id 序列转成 Tensor，用DataLoader分 batch。天池数据量大，batch_size设 64 或 128 比较稳。太小训练慢，太大显存不够且梯度更新次数少。

from torch.utils.data import Dataset, DataLoader class NewsDataset(Dataset): def __init__(self, ids, labels=None): self.ids = ids self.labels = labels def __len__(self): return len(self.ids) def __getitem__(self, idx): ids = torch.tensor(self.ids[idx], dtype=torch.long) if self.labels is not None: label = torch.tensor(self.labels[idx], dtype=torch.long) return ids, label return ids train_dataset = NewsDataset(train_df['ids'].tolist(), train_df['label'].tolist()) train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=2)

num_workers=2在 Linux 上能加速数据读取，Windows 上如果报错就设 0。shuffle=True只在训练集开，验证集和测试集不要开。

4.2 优化器、学习率和损失函数的选择

优化器我一般用 Adam，学习率 1e-3，配合ReduceLROnPlateau在验证集 loss 不降时减半。损失函数用CrossEntropyLoss，如果类别不均衡可以加weight参数。天池这个赛题类别基本均衡，不加权重也行。

import torch.optim as optim device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = LSTMClassifier(vocab_size=len(vocab)).to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=1e-3) scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=2) for epoch in range(10): model.train() total_loss = 0 for ids, labels in train_loader: ids, labels = ids.to(device), labels.to(device) optimizer.zero_grad() logits = model(ids) loss = criterion(logits, labels) loss.backward() # 梯度裁剪，防止 LSTM 梯度爆炸 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0) optimizer.step() total_loss += loss.item() print(f'Epoch {epoch}, Loss: {total_loss / len(train_loader):.4f}')

clip_grad_norm_的max_norm=5.0是 LSTM 训练的后悔药，不加的话 loss 可能突然变 NaN。patience=2表示验证集 loss 连续 2 个 epoch 不降就减学习率。如果显存够，batch_size可以加到 256，学习率相应调到 2e-3。

4.3 验证集划分和早停策略

训练集不能全用来训练，要切 10% 做验证。早停是防止过拟合最直接的手段，验证集 loss 连续 3 个 epoch 不降就停。

from sklearn.model_selection import train_test_split train_ids, val_ids, train_labels, val_labels = train_test_split( train_df['ids'].tolist(), train_df['label'].tolist(), test_size=0.1, random_state=42) val_dataset = NewsDataset(val_ids, val_labels) val_loader = DataLoader(val_dataset, batch_size=128, shuffle=False) best_val_loss = float('inf') patience_counter = 0 for epoch in range(20): # 训练部分省略，同上 model.eval() val_loss = 0 with torch.no_grad(): for ids, labels in val_loader: ids, labels = ids.to(device), labels.to(device) logits = model(ids) val_loss += criterion(logits, labels).item() val_loss /= len(val_loader) scheduler.step(val_loss) if val_loss < best_val_loss: best_val_loss = val_loss torch.save(model.state_dict(), 'best_model.pt') patience_counter = 0 else: patience_counter += 1 if patience_counter >= 3: print('早停触发') break

random_state=42保证每次切分一致，方便对比实验。保存best_model.pt而不是最后一个 epoch 的模型，因为最后一个可能已经过拟合。早停的patience设 3 比较稳，设 1 容易误停。

5. 避坑与排查：LSTM 新闻分类里最容易翻车的 5 个地方

5.1 现象：loss 一直是 NaN，训练无法继续

原因通常是学习率太大或梯度爆炸。LSTM 的梯度回传路径长，不加裁剪很容易爆。解决：把学习率降到 1e-4，加上clip_grad_norm_(model.parameters(), max_norm=5.0)。如果还是 NaN，检查输入 id 里有没有超出词表大小的值，nn.Embedding遇到越界 id 会直接报错或产生 NaN。

5.2 现象：训练集准确率 99%，验证集只有 60%

这是典型过拟合。原因可能是模型太大、训练轮数太多、dropout 太小。解决：把hidden_dim从 256 降到 128，dropout从 0.3 提到 0.5，加早停。另外检查验证集是不是和训练集同分布，天池的测试集和训练集分布基本一致，如果验证集掉点严重，多半是过拟合而不是分布问题。

5.3 现象：报错RuntimeError: Expected hidden[0] size (2, 128, 256), got (2, 128, 128)

这是双向 LSTM 的 hidden 维度没对齐。hidden_dim设 256，双向后输出 512，但初始化 hidden 时如果按 256 写就会报这个错。解决：不手动初始化 hidden，让 LSTM 自己初始化，或者把 hidden 的维度写成(num_layers * 2, batch, hidden_dim)。我一般直接不传 hidden，省事。

5.4 现象：词表建完发现<UNK>占比超过 20%

说明词表太小或者过滤太狠。原因可能是count < 2把太多词过滤了，或者max_len截断导致长尾词没统计到。解决：把词表上限提到 20 万，count < 2改成count < 1，即保留所有词。但词表太大会让 Embedding 参数量暴涨，需要权衡。另一个办法是用字符级 token，但新闻分类里词级通常更好。

5.5 现象：测试集提交后分数比验证集低很多

原因可能是测试集里有些词没在词表里，全变<UNK>。解决：建词表时合并测试集文本，确保测试集里的词也在词表中。另外检查测试集的 id 转换函数和训练集是否完全一致，max_len和截断策略要相同。如果验证集用了pack_padded_sequence而测试集没用，也会导致不一致。

6. 进阶技巧：用预训练词向量和 FGM 对抗训练再提 2 个点

6.1 加载预训练词向量初始化 Embedding

随机初始化的 Embedding 在 20 万条数据上能学到不错的表示，但如果有预训练词向量，比如腾讯词向量或 Word2Vec，加载后冻结或微调，通常能再提 1 到 2 个点。加载逻辑是遍历词表，找到对应词向量填进 Embedding 矩阵。

import numpy as np def load_pretrained_embedding(vocab, embed_dim=300): # 假设预训练词向量存为 dict: word -> np.array pretrained = {} with open('word_vectors.txt', 'r', encoding='utf-8') as f: for line in f: parts = line.strip().split() word = parts[0] vec = np.array([float(x) for x in parts[1:]]) pretrained[word] = vec embedding_matrix = np.random.normal(0, 0.1, (len(vocab), embed_dim)) hit = 0 for word, idx in vocab.items(): if word in pretrained: embedding_matrix[idx] = pretrained[word] hit += 1 print(f'命中预训练词: {hit}/{len(vocab)}') return embedding_matrix # 赋值给 Embedding 层 embedding_matrix = load_pretrained_embedding(vocab) model.embedding.weight.data.copy_(torch.tensor(embedding_matrix, dtype=torch.float)) # 可以选择冻结 embedding # model.embedding.weight.requires_grad = False

命中率低于 50% 说明预训练词向量和当前语料领域差异大，效果可能不明显。冻结 Embedding 适合数据量小的情况，数据量大时微调更好。

6.2 FGM 对抗训练在 LSTM 上的实现

FGM 是一种对抗训练方法，通过在 Embedding 层加扰动，让模型对噪声更鲁棒。实现上是在每次loss.backward()后，对 Embedding 权重加一个梯度方向的扰动，再算一次 loss 并更新。

class FGM: def __init__(self, model, epsilon=1.0): self.model = model self.epsilon = epsilon self.backup = {} def attack(self): for name, param in self.model.named_parameters(): if param.requires_grad and 'embedding' in name: self.backup[name] = param.data.clone() norm = torch.norm(param.grad) if norm != 0: r_at = self.epsilon * param.grad / norm param.data.add_(r_at) def restore(self): for name, param in self.model.named_parameters(): if name in self.backup: param.data = self.backup[name] self.backup = {} # 训练循环里使用 fgm = FGM(model, epsilon=1.0) for ids, labels in train_loader: optimizer.zero_grad() logits = model(ids) loss = criterion(logits, labels) loss.backward() fgm.attack() # 加扰动 logits_adv = model(ids) loss_adv = criterion(logits_adv, labels) loss_adv.backward() fgm.restore() # 恢复 optimizer.step()

epsilon=1.0是扰动幅度，太大训练不稳定，太小没效果。FGM 只对 Embedding 层加扰动，因为 LSTM 层加扰动计算量大且收益低。加了 FGM 后训练时间增加约 30%，但验证集通常能提 0.5 到 1 个点。

6.3 验证提分是否真实的交叉验证方法

单次验证集划分有随机性，可能这次提了下次没提。我一般用 5 折交叉验证，每折都跑一遍，看平均分和方差。如果 FGM 在 5 折里平均提了 0.8 个点，方差不大，那才是真提分。

from sklearn.model_selection import KFold kf = KFold(n_splits=5, shuffle=True, random_state=42) scores = [] for fold, (tr_idx, va_idx) in enumerate(kf.split(train_df)): tr_ids = train_df['ids'].iloc[tr_idx].tolist() tr_labels = train_df['label'].iloc[tr_idx].tolist() va_ids = train_df['ids'].iloc[va_idx].tolist() va_labels = train_df['label'].iloc[va_idx].tolist() # 训练模型并评估，记录验证集准确率 # scores.append(acc) print(f'5 折平均准确率: {np.mean(scores):.4f}, 标准差: {np.std(scores):.4f}')

标准差超过 0.5 个点说明模型不稳定，需要检查数据划分或模型初始化。我自己的习惯是每次改完模型先跑 5 折，确认稳定后再上全量数据训练提交。这套流程跑下来，LSTM 在天池新闻文本分类上做到 0.92 到 0.94 的准确率是正常的，再往上就要换 BERT 了。希望帮到你。

本文还有配套的精品资源，点击获取