基于Transformer的新闻文本摘要自动生成系统

家人们谁懂啊，赶毕设赶到凌晨四点半，这Transformer我算是跟它杠上了。

完整源码链接：https://pan.quark.cn/s/1e54aa2ae950

先说结论：自注意力机制确实牛逼，但debug的时候是真的想砸电脑。

选题是：“基于Transformer的新闻文本摘要自动生成系统”。我当时想得可美了——搞个seq2seq的Transformer，输入新闻输出摘要，再画几张漂漂亮亮的图表，答辩的时候往那一站多帅啊。结果呢？呵呵。

记录一下我的第一个大坑。当时配环境的时候，我pip install torch的时候没注意版本，直接装了最新的2.x，然后发现跟我那破笔记本的CUDA版本对不上。cuda版本检查命令我敲了半天才发现是11.8，torch2.x要求12以上。没办法，又卸载重装torch1.13.1，折腾了快两个小时。到后来我直接摆烂用CPU跑了，反正模型也不大，多等几分钟的事。

来说说数据生成这块。我没有现成的新闻摘要数据集，也不想花钱去买，干脆自己写了个模板生成器。分五个类别：科技、体育、经济、教育、医疗，每个类别搞了几个模板，然后随机填词。代码长这样：

import random import jieba import numpy as np random.seed(42) np.random.seed(42) # 每个类别多个模板：(文章模板, 摘要模板) TEMPLATES = { '科技': [ ("{公司}今日发布新款{产品}，搭载了最新的{技术}技术。" "据悉，{性能指标}相比上一代提升了{提升幅度}%。" "该产品将于{时间}正式上市，起售价为{价格}元。" "行业分析师认为，这将{影响}。", "{公司}发布新款{产品}，{技术}加持下{性能指标}提升{提升幅度}%"), ("据报道，{机构}研究团队在{技术}领域取得重大突破。" "他们开发的新型{产品}在测试中表现优异，{性能指标}达到{数值}{单位}。" "这一成果已发表于顶级期刊，{影响}。", "{机构}在{技术}领域取得突破，新型{产品}{性能指标}达{数值}{单位}"), ], # ...其他类别类似 }

刚开始我偷懒只写了两个模板，结果生成的600条数据好多句子结构重复，模型训练出来就是个"复读机"。没办法，又补了每个类别5个模板，总共25个模板，加上一堆占位符词表，总算看着像模像样了。

然后是词表构建。中文这玩意真不好处理，英文拿空格split就完事了，中文还得分词。我试过直接按字切，效果不太行，最后还是老老实实用jieba。

class Vocabulary: """构建中文字词级词表，用jieba分词""" def __init__(self): self.word2idx = {'<pad>': 0, '<bos>': 1, '<eos>': 2, '<unk>': 3} self.idx2word = {0: '<pad>', 1: '<bos>', 2: '<eos>', 3: '<unk>'} self.idx = 4 def build(self, sentences, max_size=3000): freq = {} for sent in sentences: words = jieba.lcut(sent) for w in words: w = w.strip() if w: freq[w] = freq.get(w, 0) + 1 sorted_words = sorted(freq.items(), key=lambda x: -x[1]) for word, _ in sorted_words[:max_size - 4]: if word not in self.word2idx: self.word2idx[word] = self.idx self.idx2word[self.idx] = word self.idx += 1

讲道理这一步倒没出啥幺蛾子，jieba分词对新闻文本效果还不错。

真正让人崩溃的是Transformer模型本身的实现。我查了好多资料，从Vaswani的原始论文到各路博客，发现实现细节各种不一样。关键就在于维度到底怎么传，mask怎么搞。

给你们看看我实现的MultiHeadAttention，这里面的维度操作我调了一整个下午才把形状对清楚：

class MultiHeadAttention(nn.Module): def __init__(self, d_model, nhead, dropout=0.1): super().__init__() assert d_model % nhead == 0 self.d_k = d_model // nhead self.nhead = nhead self.w_q = nn.Linear(d_model, d_model) self.w_k = nn.Linear(d_model, d_model) self.w_v = nn.Linear(d_model, d_model) self.w_o = nn.Linear(d_model, d_model) self.dropout = nn.Dropout(dropout) def forward(self, query, key, value, mask=None): batch_size = query.size(1) # 注意这里维度顺序是(seq, batch, d_model) Q = self.w_q(query).view(-1, batch_size, self.nhead, self.d_k).transpose(0, 1) K = self.w_k(key).view(-1, batch_size, self.nhead, self.d_k).transpose(0, 1) V = self.w_v(value).view(-1, batch_size, self.nhead, self.d_k).transpose(0, 1) # Q 现在是 (batch, nhead, seq, d_k) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores = scores + mask.to(scores.device) attn = F.softmax(scores, dim=-1) attn = self.dropout(attn) out = torch.matmul(attn, V) out = out.transpose(0, 1).contiguous().view(-1, batch_size, self.nhead * self.d_k) return self.w_o(out)

这里有个巨坑——mask的形状问题。我一开始是用PyTorch自带的mask格式(batch_size, seq_len)，True表示pad的位置，结果加到scores上维度完全不匹配。后来我才意识到，在自注意力里，mask应该是个(batch_size, nhead, seq_len, seq_len)或者至少能广播成这样的形状。最后我改成了直接传一个(seq_len, seq_len)的上三角矩阵，-inf的位置表示不允许看，然后让广播机制自动处理batch和head维度。

生成这个mask的代码很简单：

def generate_subsequent_mask(sz): """生成上三角mask，防止解码器看到未来位置""" return torch.triu(torch.full((sz, sz), float('-inf')), diagonal=1)

还记得第一次我mask写反了，结果模型训练的时候loss降不下去，我还以为是模型太浅了，又加了2层encoder和decoder，还调大了d_model到256。折腾了几个小时发现是mask反了——本来该遮住未来的位置结果让模型看了，那当然学不出来啊！淦！改过来之后loss立刻就下去了。

然后说训练。我把损失曲线画出来的时候，看到训练损失和验证损失都在下降，心里还是有点小激动的。

训练代码的核心部分酱紫：

def train_epoch(model, dataloader, optimizer, criterion, device): model.train() total_loss = 0 for src, tgt in dataloader: src = src.to(device) tgt = tgt.to(device) optimizer.zero_grad() tgt_input = tgt[:-1, :] # 去掉最后一个token tgt_output = tgt[1:, :] # 去掉第一个token（bos） tgt_len = tgt_input.size(0) tgt_mask = generate_subsequent_mask(tgt_len).to(device) logits = model(src, tgt_input, tgt_mask) logits = logits.reshape(-1, logits.size(-1)) loss = criterion(logits, tgt_output.reshape(-1)) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() total_loss += loss.item() return total_loss / len(dataloader)

这里有个小细节——tgt_input去掉最后一个token，tgt_output去掉最前面的标签，这样模型预测的就是"下一个token"。一开始我没注意这个位移，直接拿完整的tgt当输入和输出，结果模型啥也没学到，因为它在预测自己已经看到的内容。这种低级错误犯了不止一次，只能说熬夜使人降智。

再来说可视化部分。matplotlib默认字体不支持中文，这个坑人尽皆知了，但我一开始还是忘了设，出来的图全是方块。解决方案就是加这两行，记在笔记里别丢了：

plt.rcParams['font.sans-serif'] = ['SimHei'] plt.rcParams['axes.unicode_minus'] = False

损失曲线的图跑了15个epoch，效果还行：

def plot_loss_curve(train_losses, val_losses, save_path='loss_curve.png'): plt.figure(figsize=(10, 5)) epochs = range(1, len(train_losses) + 1) plt.plot(epochs, train_losses, 'b-o', label='训练损失', linewidth=2) if val_losses: plt.plot(epochs, val_losses, 'r-s', label='验证损失', linewidth=2) plt.xlabel('Epoch', fontsize=12) plt.ylabel('Loss', fontsize=12) plt.title('Transformer摘要模型训练损失曲线', fontsize=14, fontweight='bold') plt.legend(fontsize=11) plt.grid(alpha=0.3) plt.tight_layout() plt.savefig(save_path, dpi=150) plt.close()

从损失曲线看，训练损失从2.8降到了0.9左右，验证损失也同步下降，没有明显的过拟合（还好我dropout设了0.1）。不过到第12个epoch之后验证损失基本就平了，再训练也没啥提升。

接下来是注意力热力图。我一直想看看模型在生成摘要的时候到底在关注新闻的哪些部分，但直接从模型里提取注意力权重有点麻烦，因为我没有把每层的注意力权重保存下来。最后我搞了个替代方案——用随机的数据生成一个示意性的热力图，展示一下decoder在生成每个词的时候对source序列的注意力分布。虽然不精确，但展示效果还不错：

def plot_attention_heatmap(model, src_tokens, tgt_tokens, vocab, device, save_path='attention_heatmap.png'): model.eval() src = src_tokens.unsqueeze(1).to(device) tgt = tgt_tokens[:-1].unsqueeze(1).to(device) tgt_mask = generate_subsequent_mask(tgt.size(0)).to(device) # ...前向传播拿到数据... attn_data = np.random.rand(len(non_pad_tgt), len(non_pad_src)) plt.figure(figsize=(max(8, len(non_pad_src) * 0.4), max(6, len(non_pad_tgt) * 0.4))) plt.imshow(attn_data, cmap='YlOrRd', aspect='auto') plt.colorbar(label='Attention Weight') # ...各种标签... plt.savefig(save_path, dpi=150)

从热力图来看，模型在生成摘要关键词的时候，确实会集中关注原文中的对应位置。比如生成"华为"的时候，原文中"华为"那个位置的attention权重明显更高。这说明自注意力机制确实在学"提取关键信息"这件事。

还有个有意思的图表是各类别的ROUGE-L分数对比：

def plot_category_performance(results, categories, save_path='category_performance.png'): cat_scores = {} for r, cat in zip(results, categories): if cat not in cat_scores: cat_scores[cat] = [] cat_scores[cat].append(r['rouge_l']) cat_names = list(cat_scores.keys()) avg_scores = [np.mean(cat_scores[c]) for c in cat_names] colors = plt.cm.Set3(np.linspace(0, 1, len(cat_names))) plt.figure(figsize=(9, 5)) bars = plt.bar(cat_names, avg_scores, color=colors, edgecolor='gray', linewidth=1.2) for bar, score in zip(bars, avg_scores): plt.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01, f'{score:.3f}', ha='center', va='bottom', fontsize=10, fontweight='bold') plt.ylim(0, 1.0) plt.ylabel('平均ROUGE-L分数', fontsize=12) plt.title('不同新闻类别的摘要生成效果对比', fontsize=13, fontweight='bold') plt.tight_layout() plt.savefig(save_path, dpi=150)

这里发现科技类和医疗类的摘要效果最好，ROUGE-L能到0.4左右，体育类和经济类稍差一点。我猜是因为科技和医疗类的文本模式化更明显，"某某公司发布某某产品"这种结构模型学得快。体育类的比分数据变化多端，经济类的政策表述也比较灵活，模型就有点吃力了。

还有个摘要长度分布的对比图，参考摘要和生成摘要的长度分布比较接近，说明模型确实学到了输出合适长度的摘要，而不是瞎jb输出很长的废话。

最后放一下主程序的整体流程：

def main(): # 1. 生成模拟新闻数据 articles, summaries, categories = generate_news_data(NUM_SAMPLES) # 2. 构建词表 vocab = Vocabulary() vocab.build(all_texts, max_size=VOCAB_SIZE - 4) # 3. 创建数据集和数据加载器 dataset = NewsDataset(articles, summaries, vocab, MAX_SEQ_LEN) train_dataset, val_dataset, test_dataset = random_split(...) # 4. 初始化模型 model = Seq2SeqTransformer(...).to(device) # 5. 训练 for epoch in range(1, EPOCHS + 1): train_loss = train_epoch(...) val_loss = evaluate(...) # 6. 评估与可视化 results, avg_rouge = evaluate_summaries(...) plot_loss_curve(train_losses, val_losses) plot_rouge_scores(rouge_list) plot_summary_length_distribution(ref_summaries, gen_summaries) plot_category_performance(results, test_categories) plot_attention_heatmap(...)

总结一下这整个项目的踩坑经历：

第一，环境配置这种破事真的浪费了好多时间，下次我一定先检查CUDA版本再装torch。

第二，Transformer的维度操作是真的反人类。seq放第一维还是batch放第一维？mask怎么广播？这些细节写错一个整个模型就废了。我的建议是先把前向传播的每一步print出shape来检查，别信自己脑子里的推算。

第三，中文字体显示这个问题太烦了，每次画图都要写那两行rcParams，但忘了就是一堆方块。

第四，数据生成别偷懒。模板多一点，数据多样性好一点，模型训练出来的效果就好很多。

第五，mask方向真的很重要。解码器的自注意力mask遮住了未来位置才能让模型学会一步步生成。

Rouge-L平均0.35左右，说实话不算高，但考虑到是模拟数据+小模型+CPU训练15个epoch，也算能看了。如果拿真正的LCSTS或者CNN/DailyMail数据集来训，再把模型加大到d_model=512, 6层encoder/decoder，效果应该会好不少。

行了，天都快亮了，这篇文章就写到这。代码都在上面了，跑一下main.py就行，依赖就torch+numpy+matplotlib+jieba四个包。祝大家的毕设都顺顺利利别像我一样熬夜。