使用PyTorch实现LSTM生成ai诗

最近学习torch的一个小demo。

什么是LSTM？

长短时记忆网络（Long Short-Term Memory，LSTM）是一种循环神经网络（RNN）的变体，旨在解决传统RNN在处理长序列时的梯度消失和梯度爆炸问题。LSTM引入了一种特殊的存储单元和门控机制，以更有效地捕捉和处理序列数据中的长期依赖关系。

通俗点说就是：LSTM是一种改进版的递归神经网络（RNN）。它的主要特点是可以记住更长时间的信息，这使得它在处理序列数据（如文本、时间序列、语音等）时非常有效。

步骤如下

数据准备

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import random
import string
import os# 数据加载和预处理
def load_data(filepath):with open(filepath, 'r', encoding='utf-8') as file:text = file.read()return textdef preprocess_text(text):text = text.lower()text = text.translate(str.maketrans('', '', string.punctuation))return textdata_path = 'poetry.txt'  # 替换为实际的诗歌数据文件路径
text = load_data(data_path)
text = preprocess_text(text)
chars = sorted(list(set(text)))
char_to_idx = {char: idx for idx, char in enumerate(chars)}
idx_to_char = {idx: char for char, idx in char_to_idx.items()}
vocab_size = len(chars)print(f"Total characters: {len(text)}")
print(f"Vocabulary size: {vocab_size}")

模型构建

定义LSTM模型：

class LSTMModel(nn.Module):def __init__(self, input_size, hidden_size, output_size, num_layers=2):super(LSTMModel, self).__init__()self.hidden_size = hidden_sizeself.num_layers = num_layersself.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)self.fc = nn.Linear(hidden_size, output_size)self.softmax = nn.LogSoftmax(dim=1)def forward(self, x, hidden):lstm_out, hidden = self.lstm(x, hidden)output = self.fc(lstm_out[:, -1, :])output = self.softmax(output)return output, hiddendef init_hidden(self, batch_size):weight = next(self.parameters()).datahidden = (weight.new(self.num_layers, batch_size, self.hidden_size).zero_(),weight.new(self.num_layers, batch_size, self.hidden_size).zero_())return hidden

训练模型

将数据转换成LSTM需要的格式：

def prepare_data(text, seq_length):inputs = []targets = []for i in range(0, len(text) - seq_length, 1):seq_in = text[i:i + seq_length]seq_out = text[i + seq_length]inputs.append([char_to_idx[char] for char in seq_in])targets.append(char_to_idx[seq_out])return inputs, targetsseq_length = 100
inputs, targets = prepare_data(text, seq_length)# Convert to tensors
inputs = torch.tensor(inputs, dtype=torch.long)
targets = torch.tensor(targets, dtype=torch.long)batch_size = 64
input_size = vocab_size
hidden_size = 256
output_size = vocab_size
num_epochs = 20
learning_rate = 0.001model = LSTMModel(input_size, hidden_size, output_size)
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)# Training loop
for epoch in range(num_epochs):h = model.init_hidden(batch_size)total_loss = 0for i in range(0, len(inputs), batch_size):x = inputs[i:i + batch_size]y = targets[i:i + batch_size]x = nn.functional.one_hot(x, num_classes=vocab_size).float()output, h = model(x, h)loss = criterion(output, y)optimizer.zero_grad()loss.backward()optimizer.step()total_loss += loss.item()print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss/len(inputs):.4f}")

生成

def generate_text(model, start_str, length=100):model.eval()with torch.no_grad():input_eval = torch.tensor([char_to_idx[char] for char in start_str], dtype=torch.long).unsqueeze(0)input_eval = nn.functional.one_hot(input_eval, num_classes=vocab_size).float()h = model.init_hidden(1)predicted_text = start_strfor _ in range(length):output, h = model(input_eval, h)prob = torch.softmax(output, dim=1).datapredicted_idx = torch.multinomial(prob, num_samples=1).item()predicted_char = idx_to_char[predicted_idx]predicted_text += predicted_charinput_eval = torch.tensor([[predicted_idx]], dtype=torch.long)input_eval = nn.functional.one_hot(input_eval, num_classes=vocab_size).float()return predicted_textstart_string = "春眠不觉晓"
generated_text = generate_text(model, start_string)
print(generated_text)

运行结果如下：

运行的肯定不好，但至少出结果了。诗歌我这边只放了几句，可以自己通过外部文件放入更多素材。

整体代码直接运行即可：

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import random
import string# 预定义一些中文诗歌数据
text = """
春眠不觉晓，处处闻啼鸟。
夜来风雨声，花落知多少。
床前明月光，疑是地上霜。
举头望明月，低头思故乡。
红豆生南国，春来发几枝。
愿君多采撷，此物最相思。
"""# 数据预处理
def preprocess_text(text):text = text.replace('\n', '')return texttext = preprocess_text(text)
chars = sorted(list(set(text)))
char_to_idx = {char: idx for idx, char in enumerate(chars)}
idx_to_char = {idx: char for char, idx in char_to_idx.items()}
vocab_size = len(chars)print(f"Total characters: {len(text)}")
print(f"Vocabulary size: {vocab_size}")class LSTMModel(nn.Module):def __init__(self, input_size, hidden_size, output_size, num_layers=2):super(LSTMModel, self).__init__()self.hidden_size = hidden_sizeself.num_layers = num_layersself.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)self.fc = nn.Linear(hidden_size, output_size)self.softmax = nn.LogSoftmax(dim=1)def forward(self, x, hidden):lstm_out, hidden = self.lstm(x, hidden)output = self.fc(lstm_out[:, -1, :])output = self.softmax(output)return output, hiddendef init_hidden(self, batch_size):weight = next(self.parameters()).datahidden = (weight.new(self.num_layers, batch_size, self.hidden_size).zero_(),weight.new(self.num_layers, batch_size, self.hidden_size).zero_())return hiddendef prepare_data(text, seq_length):inputs = []targets = []for i in range(0, len(text) - seq_length, 1):seq_in = text[i:i + seq_length]seq_out = text[i + seq_length]inputs.append([char_to_idx[char] for char in seq_in])targets.append(char_to_idx[seq_out])return inputs, targetsseq_length = 10
inputs, targets = prepare_data(text, seq_length)# Convert to tensors
inputs = torch.tensor(inputs, dtype=torch.long)
targets = torch.tensor(targets, dtype=torch.long)batch_size = 64
input_size = vocab_size
hidden_size = 256
output_size = vocab_size
num_epochs = 50
learning_rate = 0.003model = LSTMModel(input_size, hidden_size, output_size)
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)# Training loop
for epoch in range(num_epochs):h = model.init_hidden(batch_size)total_loss = 0for i in range(0, len(inputs), batch_size):x = inputs[i:i + batch_size]y = targets[i:i + batch_size]if x.size(0) != batch_size:continuex = nn.functional.one_hot(x, num_classes=vocab_size).float()output, h = model(x, h)loss = criterion(output, y)optimizer.zero_grad()loss.backward()optimizer.step()total_loss += loss.item()print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {total_loss / len(inputs):.4f}")def generate_text(model, start_str, length=100):model.eval()with torch.no_grad():input_eval = torch.tensor([char_to_idx[char] for char in start_str], dtype=torch.long).unsqueeze(0)input_eval = nn.functional.one_hot(input_eval, num_classes=vocab_size).float()h = model.init_hidden(1)predicted_text = start_strfor _ in range(length):output, h = model(input_eval, h)prob = torch.softmax(output, dim=1).datapredicted_idx = torch.multinomial(prob, num_samples=1).item()predicted_char = idx_to_char[predicted_idx]predicted_text += predicted_charinput_eval = torch.tensor([[predicted_idx]], dtype=torch.long)input_eval = nn.functional.one_hot(input_eval, num_classes=vocab_size).float()return predicted_textstart_string = "春眠不觉晓"
generated_text = generate_text(model, start_string, length=100)
print(generated_text)