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

最新推荐文章于 2026-03-13 00:51:13 发布

原创

最新推荐文章于 2026-03-13 00:51:13 发布 · 1.6k 阅读

标签

#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__()
        sel