本文将介绍如何使用LSTM训练一个能够创作诗歌的模型。为了训练出效果优秀的模型，我整理了来自网络的4万首诗歌数据集。我们的模型可以直接使用预先训练好的参数，这意味着您无需从头开始训练，即可在自己的电脑上体验AI作诗的乐趣。我已经为您准备好了这些训练好的参数，让您能够轻松地在自己的设备上开始创作。本文将详细讲解如何在个人电脑上运行该模型，即使您没有机器学习方面的背景知识，也能轻松驾驭，让您的AI模型在自己的电脑上运行起来，体验AI创作诗歌的乐趣.所有的代码和资料都在仓库：https://gitee.com/yw18791995155/generate_poetry.git
秋风吹拂，窗外的树叶似灵动的舞者翩翩而舞，落日余晖将天际晕染成一片醉人的橘红。
与此同时，AI 于知识的瀚海中遨游，遍览数千篇文章后，开启了它的首次创作之旅。
在对近 4 万首唐诗深度学习之后，赋诗如下：
此诗颇具韵味，实乃勤勉研习之硕果。汲取全唐诗之精华，方成就这般非凡之能，常人岂易企及？
本博客将简要分析其中的技术细节，若有阐释未尽之处，在此诚挚欢迎诸君于评论区畅所欲言，各抒己见。先呈上仓库链接https://gitee.com/yw18791995155/generate_poetry.git
若诸位无暇详阅，不妨为该项目点亮 star 或进行 fork，诸君的每一份支持都将如熠熠星光，化作我砥砺前行之强劲动力源泉。言归正传，让我们一同开启打造AI诗人的旅程吧
01 环境配置
在开始之前，确保你的电脑已经安装了必要的依赖库：PyTorch 和 NumPy。安装命令如下：
pip install torch torchvision torchaudio numpy
一切就绪，我们可以开始了！
02 初识LSTM
长短期记忆网络 LSTM是一种特殊类型的循环神经网络（RNN），它被设计用来解决传统RNN在处理长序列数据时遇到的长期依赖性问题（梯度消失和梯度爆炸问题）。
LSTM的核心优势在于其能够学习并记住长期的信息依赖关系。这种能力使得LSTM在处理长文本内容时比普通RNN更为出色。LSTM网络中包含了四个主要的组件，它们通过门控机制来控制信息的流动：
- 遗忘门（Forget Gate）：决定哪些信息应该被遗忘，不再保留在单元状态中。
- 输入门（Input Gate）：决定哪些新信息将被存储在单元状态中。
- 单元状态（Cell State）：携带数据穿越时间的信息带，可以看作是LSTM的“记忆”。
- 输出门（Output Gate）：决定哪些信息将从单元状态输出到下一个隐藏状态。
这些门控机制使得LSTM能够有选择性地保留或遗忘信息，从而有效地捕捉和利用长期依赖性。这种设计灵感来源于对传统RNN在处理长序列时遗忘信息的挑战的回应，LSTM通过这些门控结构，使得网络能够更加灵活地处理时间序列数据。
03处理数据
接下来，首先要做的就是读取准备好的诗歌数据。然后对数据进行清洗，剔除那些包含特殊字符或长度不符合要求的诗歌。清洗完数据后，我们会为每首诗加上开始和结束的标志，确保生成的诗歌有明确的起止符号。
然后，我们会构建词典，为每个词分配一个唯一的索引，同时建立词汇到索引、索引到词汇的映射关系。最后，把每首诗转换成数字序列，这样就能让模型进行处理了。
import collections
import numpy as np
import torch# 定义起始和结束标记
start_token = 'B'
end_token = 'E'def process_poems(file_name):"""处理诗歌文件，将诗歌转换为数字序列，并构建词汇表。:param file_name: 诗歌文件的路径:return:- poems_vector: 诗歌的数字序列列表- word_to_idx: 词汇到索引的映射字典- idx_to_word: 索引到词汇的映射列表"""# 初始化诗歌列表poems = []# 读取文件并处理每一行with open(file_name, "r", encoding='utf-8') as f:for line in f.readlines():try:# 分割标题和内容title, content = line.strip().split(':')content = content.replace(' ', '')# 过滤掉包含特殊字符的诗歌if '_' in content or '(' in content or '（' in content or '《' in content or '[' in content or \start_token in content or end_token in content:continue# 过滤掉长度不符合要求的诗歌if len(content) < 5 or len(content) > 79:continue# 添加起始和结束标记content = start_token + content + end_tokenpoems.append(content)except ValueError as e:pass# 统计所有单词的频率all_words = [word for poem in poems for word in poem]counter = collections.Counter(all_words)words = sorted(counter.keys(), key=lambda x: counter[x], reverse=True)# 添加空格作为填充符words.append(' ')words_length = len(words)# 构建词汇到索引和索引到词汇的映射word_to_idx = {word: i for i, word in enumerate(words)}idx_to_word = [word for word in words]# 将诗歌转换为数字序列poems_vector = [[word_to_idx[word] for word in poem] for poem in poems]return poems_vector, word_to_idx, idx_to_worddef generate_batch(batch_size, poems_vec, word_to_int):"""生成批量训练数据。:param batch_size: 批量大小:param poems_vec: 诗歌的数字序列列表:param word_to_int: 词汇到索引的映射字典:return:- x_batches: 输入数据批次- y_batches: 目标数据批次"""# 计算可以生成的批次数num_example = len(poems_vec) // batch_sizex_batches = []y_batches = []for i in range(num_example):start_index = i * batch_sizeend_index = start_index + batch_size# 获取当前批次的诗歌batches = poems_vec[start_index:end_index]# 找到当前批次中最长的诗歌长度length = max(map(len, batches))# 初始化输入数据，使用空格进行填充x_data = np.full((batch_size, length), word_to_int[' '], np.int32)# 填充输入数据for row, batch in enumerate(batches):x_data[row, :len(batch)] = batch# 创建目标数据，目标数据是输入数据向右移一位y_data = np.copy(x_data)y_data[:, :-1] = x_data[:, 1:]"""x_data y_data[6,2,4,6,9] [2,4,6,9,9][1,4,2,8,5] [4,2,8,5,5]"""# 将当前批次的数据添加到列表中yield torch.tensor(x_data), torch.tensor(y_data)04创建模型
现在是时候搭建我们的 LSTM 模型了！我们将创建一个双层 LSTM 网络。双层 LSTM 比单层的更有能力捕捉复杂的模式和结构，能够更好地处理诗歌这种带有丰富语言特征的任务。
import torch
import torch.nn as nn
import torch.optim as optimclass RNNModel(nn.Module):def __init__(self, vocab_size, rnn_size=128, num_layers=2):"""构建RNN序列到序列模型。:param vocab_size: 词汇表大小:param rnn_size: RNN隐藏层大小:param num_layers: RNN层数"""super(RNNModel, self).__init__()# 选择LSTM单元# 参数说明：输入大小、隐藏层大小、层数、batch_first=True表示输入数据的第一维是批次大小self.cell = nn.LSTM(rnn_size, rnn_size, num_layers, batch_first=True)# 嵌入层，将词汇表中的词转换为向量# vocab_size + 1 是因为在词嵌入中需要有一个特殊标记，用于表示填充位置，所以词嵌入时会加一个词。self.embedding = nn.Embedding(vocab_size + 1, rnn_size)# RNN隐藏层大小self.rnn_size = rnn_size# 全连接层，用于输出预测# 输入大小为RNN隐藏层大小，输出大小为词汇表大小加1self.fc = nn.Linear(rnn_size, vocab_size + 1)def forward(self, input_data, hidden):"""前向传播:param input_data: 输入数据，形状为 (batch_size, sequence_length):param output_data: 输出数据（训练时提供），形状为 (batch_size, sequence_length):return: 输出结果或损失"""# 获取批次大小batch_size = input_data.size(0)# 嵌入层，将输入数据转换为向量# 输入数据形状为 (batch_size, sequence_length)，嵌入后形状为 (batch_size, sequence_length, rnn_size)embedded = self.embedding(input_data)# 通过RNN层# 输入形状为 (batch_size, sequence_length, rnn_size)，输出形状为 (batch_size, sequence_length, rnn_size)outputs, hidden = self.cell(embedded, hidden)# 将输出展平# 展平后的形状为 (batch_size * sequence_length, rnn_size)outputs = outputs.contiguous().view(-1, self.rnn_size)# 通过全连接层# 输入形状为 (batch_size * sequence_length, rnn_size)，输出形状为 (batch_size * sequence_length, vocab_size + 1)logits = self.fc(outputs)return logits, hidden
05训练模型
接下来，就是我们最考验耐性的部分——训练模型了。训练过程中，你可能需要一些时间，所以建议使用 GPU 加速。经过实测，使用 GPU 训练速度大约是 CPU 的四倍左右。所以，如果你有条件，最好让 GPU 出马，省时省力。
import torch
from model import RNNModel
from torch import nn
from poem_data_processing import *
import os
import time# 检查是否有可用的GPU，如果没有则使用CPU
# windows用户使用torch.cuda.is_available()来检查是否有可用的GPU。
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")print(f"Using device: {device}")def train(poems_path, num_epochs, batch_size, lr):"""训练RNN模型并进行预测。参数:poems_path (str): 诗歌数据文件路径。num_epochs (int): 训练的轮数。batch_size (int): 批次大小。lr (float): 学习率。"""# 确保模型保存目录存在if not os.path.exists('./model'):os.makedirs('./model')# 处理诗歌数据，生成向量化表示和映射字典poems_vector, word_to_idx, idx_to_word = process_poems(poems_path)# 初始化RNN模型并将其移动到指定设备model = RNNModel(len(idx_to_word), 128, num_layers=2).to(device)# 使用Adam优化器初始化训练器trainer = torch.optim.Adam(model.parameters(), lr=lr)# 使用交叉熵损失函数loss_fn = nn.CrossEntropyLoss()# 开始训练过程for epoch in range(num_epochs):loss_sum = 0start = time.time()# 生成并迭代训练批次for X, Y in generate_batch(batch_size, poems_vector, word_to_idx):# 将输入和目标数据移动到指定设备X = X.to(device)Y = Y.to(device)state = None# 前向传播outputs, state = model(X, state)Y = Y.view(-1)# 计算损失l = loss_fn(outputs, Y.long())# 反向传播和优化trainer.zero_grad()l.backward()torch.nn.utils.clip_grad_norm_(model.parameters(), 0.01)trainer.step()loss_sum += l.item() * Y.shape[0]end = time.time()print(f"Time cost: {end - start}s")print(f"epoch: {epoch}, loss: {loss_sum / len(poems_vector)}")# 保存模型和优化器的状态try:torch.save({'model_state_dict': model.state_dict(),'optimizer_state_dict': trainer.state_dict(),}, os.path.join('./model', 'torch-latest.pth'))except Exception as e:print(f"Error saving model: {e}")if __name__ == "__main__":file_path = "./data/poems.txt"train(file_path, num_epochs=100, batch_size=64, lr=0.002)经过了约 20 分钟的训练，终于，模型训练完成！训练结束后，模型的参数会自动保存到文件中，这样下次就可以直接加载预训练的模型，省去重新训练的麻烦。
06测试模型
终于，我们来到了最激动人心的环节——AI作诗。经过几个小时的努力，我们的AI诗人已经准备好创作一首藏头诗，以此来弥补我因编程而失去的头发。
鸡枝蝉及九层峰，内邸曾随佛统衣。
你写明时何处寻，大江蕃戴帝来儿。
太古能弗岂何如，惟无百物恣蹉跎。
美人迟意识王机，马首辞来六堕愁。
测试代码
# 导入必要的库
import torch
from model import RNNModel
from poem_data_processing import process_poems
import numpy as np# 定义开始和结束标记
start_token = 'B'
end_token = 'E'
# 模型保存的目录
model_dir = './model/'
# 诗歌数据文件路径
poems_file = './data/poems.txt'# 学习率
lr = 0.0002def to_word(predict, vocabs):"""将预测结果转换为词汇表中的字。参数:predict: 模型的预测结果，一个概率分布。vocabs: 词汇表，包含所有可能的字。返回:从预测结果中随机选择的一个字。"""predict = predict.numpy()[0]predict /= np.sum(predict)sample = np.random.choice(np.arange(len(predict)), p=predict)if sample > len(vocabs):return vocabs[-1]else:return vocabs[sample]def gen_poem(begin_word):"""生成诗歌。参数:begin_word: 诗歌的第一个字。返回:生成的诗歌，以字符串形式返回。"""batch_size = 1# 处理诗歌数据，得到诗歌向量、字到索引的映射和索引到字的映射poems_vector, word_to_idx, idx_to_word = process_poems(poems_file)# 初始化模型model = RNNModel(len(idx_to_word), 128, num_layers=2)# 加载模型参数checkpoint = torch.load(f'{model_dir}/torch-latest.pth')model.load_state_dict(checkpoint['model_state_dict'], strict=False)model.eval()# 初始化输入序列x = torch.tensor([word_to_idx[start_token]], dtype=torch.long).view(1, 1)hidden = None# 生成诗歌with torch.no_grad():output, hidden = model(x, hidden)predict = torch.softmax(output, dim=1)word = begin_word or to_word(predict, idx_to_word)poem_ = ''i = 0while word != end_token:poem_ += wordi += 1if i > 24:breakx = torch.tensor([word_to_idx[word]], dtype=torch.long).view(1, 1)output, hidden = model(x, hidden)predict = torch.softmax(output, dim=1)word = to_word(predict, idx_to_word)return poem_def pretty_print_poem(poem_):"""格式化打印诗歌。参数:poem_: 生成的诗歌，以字符串形式输入。"""poem_sentences = poem_.split('。')for s in poem_sentences:if s != '' and len(s) > 10:print(s + '。')if __name__ == '__main__':# 用户输入第一个字begin_char = input('请输入第一个字 please input the first character: \n')print('AI作诗 generating poem...')# 生成诗歌poem = gen_poem(begin_char)# 打印诗歌pretty_print_poem(poem_=poem)效果出乎意料地好，所有的努力都值了。是不是觉得很有趣？快来下载代码，亲自体验AI作诗的乐趣吧。
项目目录
- 训练模型，运行train.py文件。
- 想直接体验AI作诗，运行test.py文件。
如果不想从头训练，可以直接使用预训练好的模型参数，这些参数已经保存在文件中，只需下载仓库的所有代码和文件即可。
仓库地址：https://gitee.com/yw18791995155/generate_poetry.git
读到这里，如果你觉得这篇文章有点意思，不妨转发点赞。如果你对AI小项目感兴趣，欢迎关注我，我会持续分享更多有趣的项目。
感谢你的阅读，愿你的代码永远没有bug，头发永远浓密！
 —— 筑梦之路)
