本文将介绍如何使用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_token
poems.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_word
def 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_size
x_batches = []
y_batches = []
for i in range(num_example):
start_index = i * batch_size
end_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 optim
class 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隐藏层大小，输出大小为词汇表大小加1
self.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 = 0
start = 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.0002
def 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 = 0
while word != end_token:
poem_ += word
i += 1
if i > 24:
break
x = 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，头发永远浓密！
