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

一只小铁柱

于 2025-03-11 15:03:56 发布

阅读量1.2k

点赞数 25

文章标签： lstm pytorch 人工智能

本文链接：https://blog.csdn.net/beautiful77moon/article/details/146178539

版权

一、代码实现

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. 各个文件的作用

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

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

model.py：定义LSTM网络结构

train.py：模型训练流程

generate.py：文本生成功能

2. 文件执行顺序

运行train.py进行模型训练

使用generate.py生成文本

可调整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：实现文本生成逻辑。