一个基于LSTM的字符级文本生成模型的训练+使用(pytorch)

一、​​​​代码实现

1. 配置文件  config.py

import torch
# 设备配置
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 超参数和配置
SEQ_LENGTH = 100       # 输入序列长度
BATCH_SIZE = 64        # 批大小
EMBEDDING_DIM = 256    # 嵌入层维度
HIDDEN_SIZE = 512      # LSTM隐藏层大小
NUM_LAYERS = 2         # LSTM层数
NUM_EPOCHS = 30        # 训练轮数
LEARNING_RATE = 0.001  # 学习率
TEXT_PATH = 'shakespeare.txt'  # 文本路径
MODEL_SAVE_PATH = 'text_generation_model.pth'  # 模型保存路径

 2. 数据加载、预处理  data_preprocessing.py

import numpy as np
import torch
from torch.utils.data import DataLoader, TensorDataset
from config import *

def preprocess_data():
    # 加载并预处理文本数据
    with open(TEXT_PATH, 'r', encoding='utf-8') as f:
        text = f.read()

    # 创建字符映射
    chars = sorted(list(set(text)))
    char_to_idx = {c: i for i, c in enumerate(chars)}
    idx_to_char = {i: c for i, c in enumerate(chars)}

    # 文本转索引序列
    text_indices = np.array([char_to_idx[c] for c in text])

    # 创建训练样本
    samples = []
    for i in range(0, len(text) - SEQ_LENGTH):
        samples.append(text_indices[i:i + SEQ_LENGTH + 1])
    samples = np.array(samples)

    # 分割输入和目标
    inputs = samples[:, :-1]
    targets = samples[:, 1:]

    # 转换为PyTorch数据集
    dataset = TensorDataset(
        torch.LongTensor(inputs),
        torch.LongTensor(targets)
    )
    
    return DataLoader(dataset, BATCH_SIZE, shuffle=True), char_to_idx, idx_to_char

 3. 模型定义   model.py

import torch.nn as nn
from config import DEVICE

class CharLSTM(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, EMBEDDING_DIM)
        self.lstm = nn.LSTM(EMBEDDING_DIM, HIDDEN_SIZE, NUM_LAYERS, batch_first=True)
        self.fc = nn.Linear(HIDDEN_SIZE, vocab_size)

        # 保存超参数用于隐藏状态初始化
        self.hidden_size = HIDDEN_SIZE
        self.num_layers = NUM_LAYERS

    def forward(self, x, hidden):
        x = self.embed(x)
        out, hidden = self.lstm(x, hidden)
        out = self.fc(out)
        return out, hidden

    def init_hidden(self, batch_size):
        """初始化隐藏状态"""
        return (torch.zeros(self.num_layers, batch_size, self.hidden_size).to(DEVICE),
                torch.zeros(self.num_layers, batch_size, self.hidden_size).to(DEVICE))

4. 训练  train.py

import torch
import torch.nn as nn
from model import CharLSTM
from data_preprocessing import preprocess_data
from config import *

def train():
    # 加载数据
    dataloader, char_to_idx, _ = preprocess_data()
    vocab_size = len(char_to_idx)

    # 初始化模型并移动到GPU
    model = CharLSTM(vocab_size).to(DEVICE)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)

    # 训练循环
    for epoch in range(NUM_EPOCHS):
        for inputs, targets in dataloader:
            # 将数据移动到GPU
            inputs = inputs.to(DEVICE)
            targets = targets.to(DEVICE)

            # 动态获取当前batch大小
            batch_size = inputs.size(0)
            hidden = model.init_hidden(batch_size)

            optimizer.zero_grad()
            outputs, hidden = model(inputs, hidden)
            loss = criterion(outputs.view(-1, vocab_size), targets.view(-1))
            loss.backward()
            optimizer.step()

            # 切断隐藏状态的历史梯度
            hidden = tuple(h.detach() for h in hidden)

        print(f'Epoch {epoch + 1}/{NUM_EPOCHS} Loss: {loss.item():.4f}')

    # 保存模型
    torch.save(model.state_dict(), MODEL_SAVE_PATH)

if __name__ == '__main__':
    train()

5. 文本生成

import torch
from model import CharLSTM
from config import *

def generate(start_seq, num_chars=500):
    # 加载字符映射
    with open(TEXT_PATH, 'r') as f:
        text = f.read()
    chars = sorted(list(set(text)))
    char_to_idx = {c: i for i, c in enumerate(chars)}
    idx_to_char = {i: c for i, c in enumerate(chars)}

    # 加载模型
    model = CharLSTM(len(chars)).to(DEVICE)
    model.load_state_dict(torch.load(MODEL_SAVE_PATH))
    model.eval()

    # 生成过程
    input_seq = torch.LongTensor([[char_to_idx[c] for c in start_seq]])
    hidden = model.init_hidden(1)  # batch_size=1

    generated = []
    for _ in range(num_chars):
        output, hidden = model(input_seq, hidden)
        prob = torch.softmax(output[0, -1], dim=-1)
        char_id = torch.multinomial(prob, 1).item()

        generated.append(idx_to_char[char_id])
        input_seq = torch.cat([input_seq[:, 1:],
                               torch.LongTensor([[char_id]]).unsqueeze(0)],
                              dim=1)

    return start_seq + ''.join(generated)

if __name__ == '__main__':
    print(generate("ROMEO:", 1000))

二、代码说明

1. 各个文件的作用 

1. config.py：集中管理所有超参数和路径配置

2. data_preprocessing.py：数据加载、预处理和数据集创建

3. model.py：定义LSTM网络结构

4. train.py：模型训练流程

5. generate.py：文本生成功能

2. 文件执行顺序

1. 运行train.py进行模型训练

2. 使用generate.py生成文本

3. 可调整config.py中的参数进行实验

3. 注意事项

• 如果GPU内存不足，可以尝试减小 BATCH_SIZE 或 HIDDEN_SIZE。

• 如果没有GPU，代码会自动回退到CPU运行。

• 可以使用 nvidia-smi 命令监控GPU的使用情况。

三、主要功能总结

1. 文本数据预处理

• 功能：从文本文件中加载数据，并将其转换为模型可处理的格式。

• 具体实现：

  • 加载文本文件。

  • 构建字符到索引的映射（char_to_idx）和索引到字符的映射（idx_to_char）。

  • 将文本转换为索引序列。

  • 将数据分割为输入序列和目标序列，并创建PyTorch的 DataLoader 以便批量加载数据。

2. LSTM模型定义

• 功能：定义一个基于LSTM的神经网络模型，用于学习文本的字符级模式。

• 具体实现：

  • 使用 nn.Embedding 将字符索引映射为稠密向量。

  • 使用 nn.LSTM 处理序列数据，捕捉文本的上下文信息。

  • 使用 nn.Linear 将LSTM的输出映射为字符概率分布。

  • 提供隐藏状态的初始化方法（init_hidden）。

3. 模型训练

• 功能：训练LSTM模型，使其能够学习文本数据的分布。

• 具体实现：

  • 使用交叉熵损失函数（nn.CrossEntropyLoss）计算损失。

  • 使用Adam优化器更新模型参数。

  • 在每个epoch中遍历数据集，计算损失并反向传播。

  • 保存训练好的模型权重到文件（text_generation_model.pth）。

4. 文本生成

• 功能：使用训练好的模型生成新的文本。

• 具体实现：

  • 加载训练好的模型权重。

  • 根据输入的起始字符串（如 "ROMEO"），逐步预测下一个字符。

  • 使用概率采样（torch.multinomial）生成多样性文本。

  • 将生成的字符拼接成完整的文本并返回。

5. 模块化设计

• 功能：将代码按功能拆分为多个模块，便于维护和扩展。

• 具体实现：

  • config.py：集中管理超参数和路径配置。

  • data_preprocessing.py：处理数据加载和预处理。

  • model.py：定义LSTM模型结构。

  • train.py：实现模型训练逻辑。

  • generate.py：实现文本生成逻辑。

6. 主要技术点

• 字符级建模：以字符为单位进行文本生成，适合生成任意文本。

• LSTM网络：使用LSTM捕捉文本的长期依赖关系。

• 概率采样：在生成文本时引入随机性，避免生成重复或单调的文本。

• 模块化设计：将数据预处理、模型定义、训练和生成逻辑分离，提高代码可读性和可维护性。

7. 应用场景

• 文本生成：生成类似莎士比亚风格的文本。

• 语言模型：学习文本数据的分布，可用于文本补全、对话生成等任务。

• 教育用途：帮助理解LSTM和字符级语言模型的工作原理。

8. 改进空间

• 性能优化：有条件的可以使用更大的数据集和更深的网络提升生成质量。

• 多样性控制：在生成过程中引入温度参数（temperature）控制生成文本的多样性。

• GPU支持：将模型和数据迁移到GPU以加速训练和生成。

• 更复杂的模型：有条件的可以尝试使用Transformer或GPT等更先进的模型架构。