基于Transformer的新闻文本摘要自动生成系统
家人们谁懂啊赶毕设赶到凌晨四点半这Transformer我算是跟它杠上了。完整源码链接https://pan.quark.cn/s/1e54aa2ae950先说结论自注意力机制确实牛逼但debug的时候是真的想砸电脑。选题是“基于Transformer的新闻文本摘要自动生成系统”。我当时想得可美了——搞个seq2seq的Transformer输入新闻输出摘要再画几张漂漂亮亮的图表答辩的时候往那一站多帅啊。结果呢呵呵。记录一下我的第一个大坑。当时配环境的时候我pip install torch的时候没注意版本直接装了最新的2.x然后发现跟我那破笔记本的CUDA版本对不上。cuda版本检查命令我敲了半天才发现是11.8torch2.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_size3000): 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(), keylambda 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, dropout0.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, maskNone): 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)), diagonal1)还记得第一次我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:, :] # 去掉第一个tokenbos 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_norm1.0) optimizer.step() total_loss loss.item() return total_loss / len(dataloader)这里有个小细节——tgt_input去掉最后一个tokentgt_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_pathloss_curve.png): plt.figure(figsize(10, 5)) epochs range(1, len(train_losses) 1) plt.plot(epochs, train_losses, b-o, label训练损失, linewidth2) if val_losses: plt.plot(epochs, val_losses, r-s, label验证损失, linewidth2) plt.xlabel(Epoch, fontsize12) plt.ylabel(Loss, fontsize12) plt.title(Transformer摘要模型训练损失曲线, fontsize14, fontweightbold) plt.legend(fontsize11) plt.grid(alpha0.3) plt.tight_layout() plt.savefig(save_path, dpi150) 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_pathattention_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, cmapYlOrRd, aspectauto) plt.colorbar(labelAttention Weight) # ...各种标签... plt.savefig(save_path, dpi150)从热力图来看模型在生成摘要关键词的时候确实会集中关注原文中的对应位置。比如生成华为的时候原文中华为那个位置的attention权重明显更高。这说明自注意力机制确实在学提取关键信息这件事。还有个有意思的图表是各类别的ROUGE-L分数对比def plot_category_performance(results, categories, save_pathcategory_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, colorcolors, edgecolorgray, linewidth1.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}, hacenter, vabottom, fontsize10, fontweightbold) plt.ylim(0, 1.0) plt.ylabel(平均ROUGE-L分数, fontsize12) plt.title(不同新闻类别的摘要生成效果对比, fontsize13, fontweightbold) plt.tight_layout() plt.savefig(save_path, dpi150)这里发现科技类和医疗类的摘要效果最好ROUGE-L能到0.4左右体育类和经济类稍差一点。我猜是因为科技和医疗类的文本模式化更明显某某公司发布某某产品这种结构模型学得快。体育类的比分数据变化多端经济类的政策表述也比较灵活模型就有点吃力了。还有个摘要长度分布的对比图参考摘要和生成摘要的长度分布比较接近说明模型确实学到了输出合适长度的摘要而不是瞎jb输出很长的废话。最后放一下主程序的整体流程def main(): # 1. 生成模拟新闻数据 articles, summaries, categories generate_news_data(NUM_SAMPLES) # 2. 构建词表 vocab Vocabulary() vocab.build(all_texts, max_sizeVOCAB_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_model512, 6层encoder/decoder效果应该会好不少。行了天都快亮了这篇文章就写到这。代码都在上面了跑一下main.py就行依赖就torchnumpymatplotlibjieba四个包。祝大家的毕设都顺顺利利别像我一样熬夜。