【LSTM文本生成器】动手写一个自动生成文章的AI，附完整代码
我是@老K玩代码，非著名IT创业者。专注分享实战项目和最新行业资讯，已累计分享超1000实战项目！
0. 前言
长短期记忆网络(即LSTM)，是一种经过优化的循环神经网络(RNN)。通过给神经元设置update gate、forget gate、output gate，有效地避免参数在长序列传递的过程中，因梯度消失而造成有效历史信息丢失的问题。
LSTM的工作原理如下：
编写成公式的话，可以写成这样的形式：
$ widetilde{c}^{< t>} = tanh(W_c[a^{< t-1>}, x^{< t>}] + b_c) $
$ Gamma_u = sigma(W_u[a^{< t-1>}, x^{< t>}] + b_u) $
$ Gamma_f = sigma(W_f[a^{< t-1>}, x^{< t>}] + b_f) $
$ Gamma_o = sigma(W_o[a^{< t-1>}, x^{< t>}] + b_o) $
$ c^{< t>} = Gamma_u * widetilde{c}^{< t>} + Gamma_f * c^{< t-1>} $
$ a^{< t>} = Gamma_o * tanh(c^{< t>}) $
关于LSTM的详细内容，建议大家可以参阅大神，或者我以往的文章。
理论知识晦涩难懂，配合实战项目学习，则会事半功倍。这里，老K分享一个有具体应用场景的项目——文本生成器，给大家一边学习一边练手。
1. 准备
开始代码前，先把需要的第三方库逐个导入项目里来：
import torch import torch.nn as nn from torch.nn.utils import clip_grad_norm_ import jieba from tqdm import tqdm
- torch就是PyTorch，我们用来搭建循环神经网络会用到的库;
- torch.nn是PyTorch下的文件，主要的模型函数都是从这个文件里获取，为了方便引用，我们把这个库文件命名成nn;
- torch.nn.utils也是PyTorch下的文件，是一些工具函数，我们这里只需要clip_grad_norm_即可;
- jieba是众所周知的中文分词工具；
- tqdm是Python自带的进度条插件工具；
2. 设计类和函数
2.1 词典映射表
我们设计一个叫Dictionary的class类，用来建议单词和索引的映射表。
class Dictionary(object): def __init__(self): self.word2idx = {} self.idx2word = {} self.idx = 0 def __len__(self): return len(self.word2idx) def add_word(self, word): if not word in self.word2idx: self.word2idx[word] = self.idx self.idx2word[self.idx] = word self.idx += 1
- __init__是这个类的初始化方法，包含了两个映射关系表：由单词映射到索引的word2idx 以及 由索引映射到单词的idx2word，以及索引指针的位置idx；
- __len__是这个类的另一个魔术方法，返回当前映射表的长度，也就是这个词典里有多少个不重复单词的数量；
- add_word是这个类最核心的方法，通过这个方法，我们可以给映射表里添加新的单词；
2.2 语料集
我们获取的语料是字符串，需要编码成计算机能运算的数值，才能进行神经网络模型的学习
所以我们设计个Corpus的class类，专门用来把文本数据数值化、向量化。
class Corpus(object): def __init__(self): self.dictionary = Dictionary() def get_data(self, path, batch_size=20): # step 1 with open(path, 'r', encoding="utf-8") as f: tokens = 0 for line in f.readlines(): words = jieba.lcut(line) + ['
'] tokens += len(words) for word in words: self.dictionary.add_word(word) # step 2 ids = torch.LongTensor(tokens) token = 0 with open(path, 'r', encoding="utf-8") as f: for line in f.readlines(): words = jieba.lcut(line) + ['
'] for word in words: ids[token] = self.dictionary.word2idx[word] token += 1 # step 3 num_batches = ids.size(0) // batch_size ids = ids[:num_batches * batch_size] ids = ids.view(batch_size, -1) return ids
- __init__是Corpus类的初始化函数，会初始化一个映射表Dictionary；
- get_data是Corpus的核心方法：
- step 1: 根据给定的path读取文件里的文本，然后遍历全部文本，把通过jieba得到的分词逐一add_word到词典映射表Dictionary；
- step 2: 实例化一个LongTensor，命名为ids。遍历全部文本，根据映射表把单词转成索引，存入ids里；
- step 3: 根据传入的batch数量batch_size，把ids重构为20行的矩阵。tensor.view是改变张量形状的方法，参数-1表示根据其它维度自动计算该维度合适的长度。
2.3 架构LSTM模型
我们会从torch.nn继承Module类，进行设置，用来训练整个循环神经网络
class LSTMmodel(nn.Module): def __init__(self, vocab_size, embed_size, hidden_size, num_layers): super(LSTMmodel, self).__init__() self.embed = nn.Embedding(vocab_size, embed_size) self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True) self.linear = nn.Linear(hidden_size, vocab_size) def forward(self, x, h): x = self.embed(x) out, (h, c) = self.lstm(x, h) out = out.reshape(out.size(0) * out.size(1), out.size(2)) out = self.linear(out) return out, (h, c)
- __init__是LSTMmodel的初始函数，依次初始了以下内容
- embed: 通过nn.Embedding初始化一个词嵌入层，用来将映射的one-hot向量词向量化。输入的参数是映射表长度(vocab_size即单词总数)和词嵌入空间的维数(embed_size即每个单词的特征数)
- lstm: 通过nn.LSTM初始化一个LSTM层，是整个模型最核心、也是唯一的隐藏层。输入的参数是词嵌入空间的维数(embed_size即每个单词的特征数)、隐藏层的节点数(即hidden_size)和隐藏层的数量(即num_layers)
- linear: 通过nn.Linear初始化一个全连接层，用来把神经网络的运算结果转化为单词的概率分布。输入的参数是LSTM隐藏层的节点数(即hidden_size)和所有单词的数量(即vocab_size)
- forward定义了这个模型的前向传播逻辑，传入的参数是输入值矩阵x和上一次运算得到的参数矩阵h：
- 用embed把输入的x词嵌入化；
- 用词嵌入化的x和上一次传递进来的参数矩阵h，对lstm进行依次迭代运算，得到输出结果out以及参数矩阵h和c；
- 将out变形(重构)为合适的矩阵形状；
- 用linear把out转为和单词一一对应的概率分布。
执行训练
有了上面的基础，我们就可以对我们的模型进行训练了
embed_size = 128 hidden_size = 1024 num_layers = 1 num_epochs = 5 batch_size = 50 seq_length = 30 learning_rate = 0.001 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
我们先设置好训练会用到的参数变量：
- embed_size: 词嵌入后的特征数；
- hidden_size: lstm中隐层的节点数；
- num_layers: lstm中的隐层数量；
- num_epochs: 全文本遍历的次数；
- batch_size: 全样本被拆分的batch组数量；
- seq_length: 获取的序列长度；
- learning_rate: 模型的学习率；
- device: 设置运算用的设备实例；
corpus = Corpus() ids = corpus.get_data('sgyy.txt', batch_size) vocab_size = len(corpus.dictionary)
接下来，我们通过Corpus的get_data方法，读取语料，并对数据进行必要的预处理
- 实例一个Corpus类;
- 用get_data方法，读取目标文件里的文本，并处理成相应的batches；
- 获得当前词典映射表的长度vocab_size（这个vocab_size在设计全连接，即单词概率分布矩阵的长度时会用到）；
model = LSTMmodel(vocab_size, embed_size, hidden_size, num_layers).to(device) cost = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
这里，我们实例了训练需要的完整结构：
- model，是模型主体LSTMmodel；
- cost，是训练的损失函数，这里我们用交叉熵损失nn.CrossEntropyLoss；
- optimizer，是训练的优化器，这里我们用Adam方法对参数进行优化。
for epoch in range(num_epochs): states = (torch.zeros(num_layers, batch_size, hidden_size).to(device), torch.zeros(num_layers, batch_size, hidden_size).to(device)) for i in tqdm(range(0, ids.size(1) - seq_length, seq_length)): inputs = ids[:, i:i+seq_length].to(device) targets = ids[:, (i+1):(i+1)+seq_length].to(device) states = [state.detach() for state in states] outputs, states = model(inputs, states) loss = cost(outputs, targets.reshape(-1)) model.zero_grad() loss.backward() clip_grad_norm_(model.parameters(), 0.5) optimizer.step()
这是主循环，呈现了训练的主体逻辑：
- states是参数矩阵的初始化，相当于对LSTMmodel类里的(h, c)的初始化；
- 在迭代器上包裹tqdm，可以打印该循环的进度条；
- inputs和targets是训练集的x和y值；
- 通过detach方法，定义参数的终点位置；
- 把inputs和states传入model，得到通过模型计算出来的outputs和更新后的states；
- 把预测值outputs和实际值targets传入cost损失函数，计算差值；
- 由于参数在反馈时，梯度默认是不断积累的，所以在这里需要通过zero_grad方法，把梯度清零以下；
- 对loss进行反向传播运算；
- 为了避免梯度爆炸的问题，用clip_grad_norm_设定参数阈值为0.5；
- 用优化器optimizer进行优化.
生成文章
当模型通过上述过程，完成训练后，我们就可以用训练过的模型，自动生成文章了。
num_samples = 300 article = str() state = (torch.zeros(num_layers, 1, hidden_size).to(device), torch.zeros(num_layers, 1, hidden_size).to(device)) prob = torch.ones(vocab_size) _input = torch.multinomial(prob, num_samples=1).unsqueeze(1).to(device)
我们先完成一些初始化的工作：
- num_samples表示生成文本的长度；
- article是字符串，作为输出文本的容器；
- state是初始化的模型参数，相当于模型中的(h, c)；
- prob对应模型中的outputs，是输入变量经过语言模型得到的输出值，相当于此时每个单词的概率分布；
- _input，出于和Python自带函数input冲突，在变量明前加下划线_，是从字典里随机抽样一个单词，作为文章开头。
for i in range(num_samples): output, state = model(_input, state) prob = output.exp() word_id = torch.multinomial(prob, num_samples=1).item() _input.fill_(word_id) word = corpus.dictionary.idx2word[word_id] word = 'n' if word == '
' else word article += word print(article)
通过主循环，实现自动生成文本的功能：
- for循环num_samples次，即可生成由num_samples个单词组成的文章；
- output、state是LSTMmodel在接收到变量_input和state后的输出值；
- prob是对上一步得到的output进行指数化，加强高概率结果的权重；
- word_id，通过torch_multinomial，以prob为权重，对结果进行加权抽样，样本数为1(即num_samples)；
- 为下一次运算作准备，通过fill_方法，把最新的结果(word_id)作为_input的值；
- 从字典映射表Dictionary里，找到当前索引(即word_id)对应的单词；
- 如果获得到的单词是特殊符号(如
，句尾符号EndOfSentence)，替换成换行符；
- 将word存到article文章容器中；
- print生成的文章，将article打印出来。
总结
通过上述方法，就可以让LSTM模型自动替我们生成一些文章文本。
以下是我以《三国演义》为语料，经过一个epoch训练后得到的模型，自动生成的文本：
夏侯渊引项城濬赵云南山。可引军将切齿韦愿往插可借张引兵哨探，—不酿得中，崩寄臣居民而立。奂，降旗转加司徒王允，便赏先主姜维所讫坐定细作，傍若无人兵迎践踏。关公陇来报为兵战为，因小疮大进张飞。 且说可怜何进，正见遂通晓孔明马超亦之孙拜而出波浪袁术。病故入献酒食这，至丙寅日。孔明曰：“之处是朱灵同心？”操曰：“张翼德等。时定军山也。吾而定乘他府，蜀兵实为也？死罪相助禳，楮并举良谋乎为即命？问时满宠精兵姜维兵，山坚守殃及，不十合坐者不满火归坐守，选长叹、曹军入从吞并，果是痛饮、护卫军、公当速、众韩关之所学门、质入彪，只得三万，跃起潘隐谓，肩同归中。鼓噪托病赵彦杀，三声已危数十字子翼，杀入破绽飞乃入，皆创立大半六年人口。左右军，皆不能使人往去 关公横截樊城。众军击班者黄门其肉都督隆冬事截杀。忽起凋残营寨。望此不到别船刺臂，今卓齐自于所舵，然后虎豹曰：“兄为何人，秋天追夺术之功，乃大魏听令经典，不忧姜维归之今蜀兵名将。今晚之精兵。”荆棘甚妙之。允曰：“吾与将军归家好以此？”遂夏侯拜谢扬妻女而。两阵徐州惰慢兵。操大惊，引曹洪领进酒具言前，不觉两军两军会小校，坛自守。操曰：“何不同在关某！”分付平：“此医与文长阴平探其防护以金帛同扶。” 黄忠孙先锋齐声见山谷，军吏小匣冬投百步成万。彧魏军曰：“贵人为红旗来！”武士膂力过人颈曰：“各引东方之心休道，献深感而相府石，使子分外将矣，难芳引路。”后人知事美髯与允并素闻密授而去。 建安荀彧改正引路。正是校尉造饭陆口守豫州动，貂蝉蒯越曰：“三处，良苦汉高祖；今不能成大功草芥，俱杀此人，安出城之辱不得、投；岸去李辅围为此如，怎敢以三人部麾抚慰矣。”孔明曰审钧意冒死慌救入引兵。布苏，壮士利斧从江众之，娱情以赐赞徐往吕旷去，班部艾。华阴羕。禳欲攻亮出受敌。偿命之，兵败将亡汉中披挂。 且说至，邓艾自大半不分昼夜至。
需要语料的可以私信我关键词 / RNN / 领取。
通过上面的例子，我们可以发现，仅仅通过一层神经网络，一轮epoch的训练，就能生成一段似是而非的文章。
我们可能可以通过以下方法进一步优化产出文本的结果：
- 调整模型内容，如将LSTM替换成GRU，或者替换损失函数和优化器；
- 增加词嵌入的特征表示embed_size，使每个单词能包含更多信息，提高计算结果的精准度；
- 提高LSTM神经元数量hidden_size或隐藏层数num_layers，以起到优化模型逻辑的作用；
- 增加训练次数，如增加num_epochs，使模型继续向最优解收敛；
- 调整增大seq_length的值，使训练传入的语句变长，增加前后词语的长距离依赖关系和准确性；
- 修改学习率learning_rate，通过不同的步长使梯度下降的过程更有效；
- 使用语法更规范，文本量更大的语料进行训练。
以上方法不一定会为模型带来更优的结果，还存在过度拟合或者其它问题的情况，各位可以根据代码，自行尝试和优化。
希望大家能基于本项目，制作出优秀的文本生成器。
作者介绍
我是@老K玩代码，非著名IT创业者。专注分享实战项目和最新行业资讯，已累计分享超1000实战项目！
全网同名，欢迎通过各种渠道和我交流。
需要源码 或者 对RNN感兴趣的小伙伴，可以私信关键词 / RNN / ，获取更多相关资料。
