从RNN到LSTM/GRU:理解递归神经网络原理与PyTorch文本生成实战
1. 从“记忆”说起:为什么需要递归神经网络?
如果你尝试用传统的全连接神经网络去预测一段文本的下一个词,比如“今天天气很”,模型会怎么处理?它会将“今天”、“天气”、“很”这几个词作为独立的特征输入,然后试图找出一个映射关系。但问题在于,模型完全不知道“今天”和“天气”这两个词在时间上的先后顺序,也不知道“很”后面大概率跟一个形容词。它丢失了序列信息,就像一个人失忆了,只记得一些零散的词语片段,却无法将它们串联成一个有意义的句子。
这就是递归神经网络(Recurrent Neural Network, RNN)要解决的核心问题:处理序列数据。在自然语言处理、语音识别、时间序列预测等领域,数据的顺序至关重要。RNN通过引入“循环”结构,让网络具备了“记忆”过去信息的能力。你可以把它想象成一个有状态的处理器,每次接收新输入时,都会结合自己当前的“记忆”(即隐藏状态)来产生输出,并更新自己的记忆,传递给下一个时刻。这种设计让RNN在处理“今天天气很”这样的序列时,能够记住“今天”和“天气”的上下文,从而更准确地预测出“好”或“坏”。
然而,经典的RNN在实践中有个致命弱点:长程依赖问题。当序列很长时(比如一段几百字的文章),早期的信息在反向传播过程中,梯度会经历连乘运算。如果这个连乘因子小于1,梯度会指数级衰减到近乎为零(梯度消失),导致网络无法学习到长距离的依赖关系;如果大于1,则会指数级爆炸(梯度爆炸),导致训练不稳定。这就像让你复述一个很长的故事,你很可能只记得最后几句,而忘记了开头的关键情节。
为了解决这个问题,后续发展出了长短期记忆网络(LSTM)和门控循环单元(GRU)等变体。它们通过引入精巧的“门控”机制,有选择性地记住重要信息、遗忘无关信息,从而有效地缓解了梯度消失问题,让网络能够学习更长的序列依赖。本文将从最基础的RNN原理讲起,逐步深入到LSTM和GRU,并最终提供一份可运行、可修改的Python训练代码,让你不仅能理解其思想,更能亲手实现它。
2. RNN的核心机制:展开计算图与时间步
理解RNN,最关键的一步是看懂它的“展开”过程。RNN的“递归”或“循环”体现在其结构参数在时间步上是共享的。
2.1 折叠与展开视图
一个基础的RNN单元,其核心是一个循环结构。在折叠视图中,它看起来像一个带自环的神经元,这个自环表示将当前时刻的输出(更准确地说,是隐藏状态)传递到下一时刻作为输入的一部分。
当我们处理一个具体的序列,比如[x(0), x(1), ..., x(t)]时,我们可以将这个循环结构按时间步展开。展开后,原本的一个RNN单元变成了t+1个结构相同的网络层,每一层对应一个时间步。关键点在于,这t+1个“层”共享同一套参数(权重W和偏置b)。这意味着,无论序列多长,我们学习的都是同一组如何从“当前输入+过去记忆”生成“当前输出和未来记忆”的规则。
用公式来表达这个核心计算过程:
隐藏状态更新:
h_t = tanh(W_{hh} * h_{t-1} + W_{xh} * x_t + b_h)h_t: 当前时刻t的隐藏状态,即网络的“记忆”。h_{t-1}: 上一时刻的隐藏状态。x_t: 当前时刻的输入。W_{hh}: 隐藏状态到隐藏状态的权重矩阵。W_{xh}: 输入到隐藏状态的权重矩阵。b_h: 隐藏层的偏置。tanh: 激活函数,将结果压缩到(-1, 1)之间,有助于稳定梯度。
输出计算:
y_t = W_{hy} * h_t + b_yy_t: 当前时刻的输出(例如,预测的下一个词的概率分布)。W_{hy}: 隐藏状态到输出的权重矩阵。b_y: 输出层的偏置。
注意:在许多任务中(如情感分类),我们可能只关心最后一个时间步的输出
y_T;而在序列生成任务中(如机器翻译),我们则需要每一个时间步的输出y_1, y_2, ..., y_T。
2.2 梯度消失与爆炸的数学根源
为什么简单的RNN会有长程依赖问题?我们来看反向传播的过程,即著名的**BPTT(Back Propagation Through Time)**算法。
假设我们在时刻t的损失函数为L_t,我们需要计算损失对更早时刻参数W_{hh}的梯度。根据链式法则,梯度∂L_t / ∂W_{hh}会包含一连乘项:∂h_t / ∂h_{t-1} * ∂h_{t-1} / ∂h_{t-2} * ... * ∂h_{k+1} / ∂h_k。
其中,∂h_{i} / ∂h_{i-1}是雅可比矩阵,其元素大小与W_{hh}和tanh的导数有关。tanh的导数范围在 (0, 1] 之间。如果W_{hh}的特征值(可以粗略理解为权重的大小)小于1,那么这一连乘项会指数级衰减到0,这就是梯度消失——早期时间步的参数几乎得不到更新。反之,如果特征值大于1,连乘项会指数级增大,导致梯度爆炸,更新步长过大,网络无法收敛。
一个直观的类比:你试图调整一个多米诺骨牌阵列中第一块骨牌的位置来影响最后一块骨牌的倒下。如果骨牌之间的传递效率(类比W_{hh})太低,力传到一半就没了(梯度消失);如果传递效率太高且不稳定,一点微调可能导致最后一块骨牌飞出去(梯度爆炸)。
3. LSTM与GRU:用“门”控制记忆流
为了解决上述问题,LSTM和GRU被提出。它们不再让信息在隐藏状态中“自由流动”,而是通过引入“门”结构来精细化控制信息的留存与遗忘。
3.1 LSTM:长短期记忆网络
LSTM在RNN的基础上,增加了一个细胞状态(Cell State),记为C_t。你可以把它看作网络的“长期记忆通道”,它贯穿整个时间线,其更新相对线性,受控于三个门。
遗忘门(Forget Gate):决定从上一细胞状态
C_{t-1}中丢弃哪些信息。- 公式:
f_t = σ(W_f · [h_{t-1}, x_t] + b_f) - 输出一个0到1之间的向量,作用在
C_{t-1}上。1表示“完全保留”,0表示“完全遗忘”。
- 公式:
输入门(Input Gate):决定将哪些新信息存入细胞状态。
- 公式:
i_t = σ(W_i · [h_{t-1}, x_t] + b_i) - 同时,一个tanh层会创建一个新的候选值向量
\tilde{C}_t = tanh(W_C · [h_{t-1}, x_t] + b_C)。
- 公式:
细胞状态更新:结合遗忘门和输入门,更新长期记忆。
- 公式:
C_t = f_t ⊙ C_{t-1} + i_t ⊙ \tilde{C}_t ⊙表示逐元素相乘。这一步是LSTM的核心:它通过加法来更新状态,而不是像简单RNN那样通过连乘。加法操作使得梯度在反向传播时以更稳定的方式流动,极大地缓解了梯度消失问题。
- 公式:
输出门(Output Gate):基于更新后的细胞状态,决定输出什么到隐藏状态
h_t。- 公式:
o_t = σ(W_o · [h_{t-1}, x_t] + b_o) h_t = o_t ⊙ tanh(C_t)
- 公式:
LSTM通过这三个门,实现了对信息流的精确控制。例如,在语言模型中,当遇到一个新的主语时,输入门可以打开,将新主语信息写入细胞状态;当一段无关的描述结束时,遗忘门可以关闭,清空相关记忆;在需要生成动词时,输出门可以从细胞状态中读取主语信息。
3.2 GRU:门控循环单元
GRU是LSTM的一个变体,它合并了细胞状态和隐藏状态,并将门数量减少到两个,结构更简单,计算效率更高,且在多数任务上表现与LSTM相当。
重置门(Reset Gate):决定如何将新的输入与之前的记忆结合。
- 公式:
r_t = σ(W_r · [h_{t-1}, x_t] + b_r)
- 公式:
更新门(Update Gate):类似于LSTM的遗忘门和输入门的结合体,决定有多少旧记忆被保留,多少新记忆被加入。
- 公式:
z_t = σ(W_z · [h_{t-1}, x_t] + b_z)
- 公式:
候选隐藏状态:结合重置门的信息,计算一个候选状态。
- 公式:
\tilde{h}_t = tanh(W · [r_t ⊙ h_{t-1}, x_t] + b) - 重置门
r_t控制着过去状态h_{t-1}对候选状态的影响。如果r_t接近0,则忽略过去状态,候选状态主要依赖于当前输入x_t,这有助于丢弃无关信息。
- 公式:
隐藏状态更新:通过更新门,在旧状态和候选状态之间进行插值。
- 公式:
h_t = (1 - z_t) ⊙ h_{t-1} + z_t ⊙ \tilde{h}_t - 更新门
z_t决定了新状态h_t中有多少来自旧状态h_{t-1}(保留),有多少来自候选状态\tilde{h}_t(更新)。当z_t接近0时,主要保留旧记忆;接近1时,主要采用新信息。
- 公式:
LSTM vs. GRU 如何选择?这是一个经验性问题。通常:
- LSTM:理论更优雅,门控机制更清晰(各有分工),在处理非常长的序列或任务对长期记忆要求极高时,有时表现更稳定。
- GRU:参数更少,训练速度更快,在大多数常见任务(如文本分类、中等长度序列预测)上能达到与LSTM相当甚至更好的效果,是当前更常用的默认选择之一。
- 实战建议:对于新任务,可以先用GRU进行快速实验和原型开发。如果效果不佳,再尝试LSTM。很多时候,数据质量、特征工程和超参数调优的影响远大于LSTM和GRU之间的选择。
4. 从零实现:一个基于PyTorch的RNN文本生成实战
理论说再多,不如亲手跑通代码。下面我们将使用PyTorch框架,实现一个基于字符级的RNN(这里以GRU为例)文本生成模型。我们的目标是让模型学习一段文本的统计规律,然后自动生成类似风格的新文本。
4.1 环境准备与数据预处理
首先,确保你的Python环境已安装PyTorch。可以使用以下命令安装(以CPU版本为例):
pip install torch torchvision torchaudio我们选择莎士比亚的戏剧文本作为训练数据。预处理步骤至关重要:
- 读取文本:将整个文本读入一个长字符串。
- 创建词汇表:找出文本中所有独特的字符,建立字符到索引(char_to_idx)和索引到字符(idx_to_char)的映射。字符级模型比词级模型词汇表小,更易于入门。
- 文本向量化:将整个文本字符串转换为对应的索引序列(Tensor)。
- 创建训练样本:我们需要将长序列切割成许多固定长度的连续片段。对于每个片段,输入是前
seq_length个字符,目标输出是后移一位的seq_length个字符(因为我们要预测下一个字符)。
import torch import torch.nn as nn import torch.optim as optim import numpy as np import random # 1. 读取数据 with open('shakespeare.txt', 'r', encoding='utf-8') as f: text = f.read() # 2. 创建词汇表 chars = sorted(list(set(text))) vocab_size = len(chars) print(f'Unique characters: {vocab_size}') print(''.join(chars[:50])) char_to_idx = {ch: i for i, ch in enumerate(chars)} idx_to_char = {i: ch for i, ch in enumerate(chars)} # 3. 文本向量化 data = torch.tensor([char_to_idx[ch] for ch in text], dtype=torch.long) # 4. 创建批次数据函数 def get_batch(split, batch_size, seq_length): # 分割训练集和验证集 data_train = data[:int(0.9*len(data))] data_val = data[int(0.9*len(data)):] data_split = data_train if split == 'train' else data_val # 随机生成批次起始点 ix = torch.randint(len(data_split) - seq_length, (batch_size,)) # 构造输入x和目标y x = torch.stack([data_split[i:i+seq_length] for i in ix]) y = torch.stack([data_split[i+1:i+seq_length+1] for i in ix]) return x, y # 测试一个批次 batch_size = 4 seq_length = 32 x_batch, y_batch = get_batch('train', batch_size, seq_length) print(f'Input batch shape: {x_batch.shape}') # [4, 32] print(f'Target batch shape: {y_batch.shape}') # [4, 32] for b in range(batch_size): print('---') print('Input:', ''.join([idx_to_char[idx.item()] for idx in x_batch[b]])) print('Target:', ''.join([idx_to_char[idx.item()] for idx in y_batch[b]]))4.2 定义GRU模型
我们将定义一个简单的GRU模型,它包含一个嵌入层(将字符索引转换为稠密向量)、一个GRU层和一个线性输出层。
class CharGRU(nn.Module): def __init__(self, vocab_size, embed_dim, hidden_size, num_layers, dropout=0.2): super().__init__() self.vocab_size = vocab_size self.hidden_size = hidden_size self.num_layers = num_layers # 嵌入层:将字符索引映射为向量 self.embedding = nn.Embedding(vocab_size, embed_dim) # GRU层:核心循环网络 self.gru = nn.GRU( input_size=embed_dim, hidden_size=hidden_size, num_layers=num_layers, batch_first=True, # 输入输出张量形状为 (batch, seq, feature) dropout=dropout if num_layers > 1 else 0 # 多层时使用dropout ) # 输出层:将GRU隐藏状态映射回词汇表空间 self.fc = nn.Linear(hidden_size, vocab_size) def forward(self, x, hidden=None): # x shape: (batch_size, seq_length) batch_size, seq_length = x.shape # 1. 通过嵌入层 embedded = self.embedding(x) # shape: (batch_size, seq_length, embed_dim) # 2. 通过GRU层 # 如果未提供初始隐藏状态,则初始化为零 if hidden is None: hidden = torch.zeros(self.num_layers, batch_size, self.hidden_size).to(x.device) gru_out, hidden_out = self.gru(embedded, hidden) # gru_out shape: (batch_size, seq_length, hidden_size) # hidden_out shape: (num_layers, batch_size, hidden_size) # 3. 通过全连接层 # 将gru_out重塑为 (batch_size * seq_length, hidden_size) 以通过全连接层 gru_out_reshaped = gru_out.contiguous().view(-1, self.hidden_size) output = self.fc(gru_out_reshaped) # shape: (batch_size * seq_length, vocab_size) # 重塑回序列形状,方便计算损失 output = output.view(batch_size, seq_length, self.vocab_size) return output, hidden_out def init_hidden(self, batch_size, device): """初始化隐藏状态,用于序列生成""" return torch.zeros(self.num_layers, batch_size, self.hidden_size).to(device)关键参数解析:
vocab_size: 词汇表大小,即独特字符的数量。embed_dim: 嵌入向量的维度。每个字符将被表示为一个embed_dim维的实数向量。这个向量会在训练中学习,相似的字符在向量空间中的位置会更接近。hidden_size: GRU隐藏状态的维度。可以理解为网络“记忆”的容量。维度越大,模型表达能力越强,但也更容易过拟合,计算量更大。num_layers: 堆叠的GRU层数。多层RNN可以学习更复杂的特征表示,但同样会增加训练难度和过拟合风险。对于字符级文本生成,1-3层通常足够。dropout: 层间丢弃率,用于防止过拟合。仅在num_layers > 1时生效。
4.3 模型训练与损失函数
我们将使用交叉熵损失函数,它非常适合多分类任务(预测下一个字符是词汇表中的哪一个)。优化器选择Adam,它是目前最常用的自适应学习率优化器。
# 超参数设置 embed_dim = 128 hidden_size = 256 num_layers = 2 dropout = 0.2 learning_rate = 0.001 batch_size = 64 seq_length = 128 num_epochs = 5000 eval_interval = 500 # 每500步评估一次 # 初始化模型、损失函数和优化器 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = CharGRU(vocab_size, embed_dim, hidden_size, num_layers, dropout).to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) # 训练循环 model.train() for epoch in range(num_epochs): # 获取一个训练批次 xb, yb = get_batch('train', batch_size, seq_length) xb, yb = xb.to(device), yb.to(device) # 前向传播 logits, _ = model(xb) # logits shape: [batch, seq, vocab] # 计算损失。需要将logits重塑为 [batch*seq, vocab], targets重塑为 [batch*seq] loss = criterion(logits.view(-1, vocab_size), yb.view(-1)) # 反向传播与优化 optimizer.zero_grad() loss.backward() # 梯度裁剪:防止梯度爆炸的关键技巧! torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() # 定期打印损失并采样生成文本 if epoch % eval_interval == 0: # 计算验证集损失 model.eval() with torch.no_grad(): xb_val, yb_val = get_batch('val', batch_size, seq_length) xb_val, yb_val = xb_val.to(device), yb_val.to(device) logits_val, _ = model(xb_val) val_loss = criterion(logits_val.view(-1, vocab_size), yb_val.view(-1)) print(f'Epoch {epoch:5d} | Train Loss: {loss.item():.4f} | Val Loss: {val_loss.item():.4f}') # 采样生成一段文本 print('--- Generated Text ---') print(generate_text(model, device, char_to_idx, idx_to_char, start_str='\n', length=200)) print('---' + '-'*20) model.train() print('Training finished.')核心技巧:梯度裁剪(Gradient Clipping)代码中的
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)是训练RNN类模型时至关重要的一步。它将所有参数的梯度拼接成一个向量,如果这个向量的范数(长度)超过了max_norm(这里设为1.0),就按比例缩放,使其范数等于max_norm。这能有效防止梯度爆炸,让训练过程更稳定。即使使用了LSTM/GRU,梯度裁剪仍然是一个良好的实践。
4.4 文本生成:采样策略
训练完成后,我们需要一个函数来让模型生成新的文本。生成过程是自回归的:给定一个起始字符串,模型预测下一个字符的概率分布,我们从这个分布中“采样”出一个字符,将其追加到输入序列末尾,然后重复这个过程。
def generate_text(model, device, char_to_idx, idx_to_char, start_str='\n', length=1000, temperature=0.8): """ 使用训练好的模型生成文本。 参数: model: 训练好的模型 device: CPU或GPU char_to_idx/idx_to_char: 词汇表映射 start_str: 生成文本的起始字符串 length: 要生成的总字符数(包括起始字符串) temperature: “温度”参数,控制生成的随机性。 temperature -> 0: 更确定,选择概率最高的字符(可能重复单调)。 temperature -> 1: 完全按概率分布采样。 temperature > 1: 更随机,探索性更强。 """ model.eval() # 将起始字符串转换为索引 chars = [ch for ch in start_str] input_seq = torch.tensor([char_to_idx[ch] for ch in chars], dtype=torch.long).unsqueeze(0).to(device) # shape: [1, len(start_str)] hidden = None with torch.no_grad(): for _ in range(length - len(start_str)): # 前向传播,获取下一个字符的logits logits, hidden = model(input_seq[:, -1:], hidden) # 只输入最后一个字符,并传入之前的隐藏状态 # logits shape: [1, 1, vocab_size] logits = logits[:, -1, :] / temperature # 应用温度参数 # 将logits转换为概率分布 probs = torch.softmax(logits, dim=-1) # shape: [1, vocab_size] # 根据概率分布采样下一个字符的索引 next_char_idx = torch.multinomial(probs, num_samples=1).item() # 将新字符索引转换为字符,并添加到序列中 chars.append(idx_to_char[next_char_idx]) # 准备下一次迭代的输入(新字符的索引) input_seq = torch.cat([input_seq, torch.tensor([[next_char_idx]], device=device)], dim=1) return ''.join(chars)温度参数(Temperature)详解: 这是控制生成文本“创造性”和“连贯性”平衡的关键旋钮。
temperature = 0.1:模型几乎总是选择概率最高的那个字符。生成结果确定性高,但可能非常重复、枯燥,缺乏多样性。temperature = 1.0:严格按模型输出的概率分布采样。这是最“忠实”于模型原始预测的方式。temperature = 1.5:概率分布被平滑,低概率字符被选中的机会增加。生成结果更具创造性、更出人意料,但也更容易出现语法错误或无意义的词。- 实战建议:通常从0.7到0.9开始尝试。如果你发现生成的文本过于保守和重复,可以适当调高(如1.0-1.2);如果生成的文本过于混乱,可以调低(如0.5-0.7)。
5. 实战中的调优策略与常见陷阱
模型跑起来只是第一步,要想获得好的生成效果,还需要在训练和生成过程中进行精细的调优。
5.1 超参数调优经验谈
- 学习率(Learning Rate):这是最重要的超参数。对于Adam优化器,
1e-3或3e-4通常是安全的起点。如果训练损失震荡剧烈或下降缓慢,可以尝试调低(如5e-4)。也可以使用学习率调度器,如ReduceLROnPlateau,在验证损失停滞时自动降低学习率。 - 批量大小(Batch Size):在GPU内存允许的范围内,使用较大的批量大小(如64, 128, 256)通常能使训练更稳定,梯度估计更准确。但批量大小也会影响泛化性能,有时小批量(如32)能带来更好的最终效果,这是一个需要权衡的点。
- 序列长度(Sequence Length):这决定了模型一次能看到的上下文长度。对于莎士比亚文本,
128或256可能足够。对于更长的依赖关系(如编程代码生成),可能需要512或更长。更长的序列意味着更深的计算图,可能加剧梯度问题,并消耗更多内存。 - 隐藏层大小与层数:更大的
hidden_size和更多的num_layers意味着模型容量更大。对于简单的字符级任务,hidden_size=256和num_layers=2是一个不错的起点。如果模型在训练集上表现很好但在验证集上很差(过拟合),可以尝试减小模型尺寸或增加dropout。 - 嵌入维度(Embedding Dim):通常设置为与
hidden_size同一数量级或略小,如128或256。它定义了字符的表示空间。
5.2 训练过程监控与调试
- 损失曲线:同时绘制训练损失和验证损失曲线。理想情况是两者都平稳下降,且验证损失最终趋于平稳。如果训练损失下降但验证损失上升,这是典型的过拟合信号,需要增加Dropout、使用权重衰减(L2正则化)、获取更多数据或进行数据增强。
- 梯度范数:可以定期打印模型参数的梯度范数。如果范数突然变得非常大(如>100),说明可能发生了梯度爆炸,需要检查梯度裁剪是否生效,或降低学习率。
- 文本生成质量:定期(如每500或1000个训练步)用固定的
start_str和temperature生成一段文本。直观地观察生成文本从随机乱码逐渐变得有语法、有结构的过程,是判断模型是否在学习的最直接方式。
5.3 从字符级到词级与更高级的架构
我们的示例是字符级模型,它优点在于词汇表小,能生成任何字符(包括标点、罕见词)。但其缺点也很明显:序列非常长,难以捕捉长距离语义依赖,训练较慢。
进阶方向:
- 词级模型:将文本按单词分割。词汇表会变大很多,需要处理未登录词(OOV)。通常需要更大的嵌入维度和模型容量,但能更有效地学习语义和句法。
- 使用预训练词向量:如GloVe或Word2Vec。用这些在大规模语料上训练好的静态词向量初始化嵌入层,可以加速收敛并提升模型性能,尤其是在训练数据有限的情况下。
- 双向RNN(Bi-RNN):对于某些任务(如序列标注、情感分析),上下文信息来自前后两个方向。Bi-RNN通过同时运行前向和后向RNN,并将两者的隐藏状态拼接,来获取更丰富的上下文信息。
- 注意力机制(Attention)与Transformer:这是当前的主流。注意力机制允许模型在生成每个输出时,直接“关注”输入序列的任何部分,彻底解决了长程依赖问题。Transformer完全基于自注意力机制,并行度高,在机器翻译、文本生成(如GPT系列)等领域取得了革命性成功。当你掌握了RNN/LSTM/GRU后,学习Transformer是自然的下一步。
6. 完整训练代码整合与运行指南
以下是整合后的完整训练脚本,你可以将其保存为train_rnn_text_generation.py并运行。
import torch import torch.nn as nn import torch.optim as optim import numpy as np import random # --- 1. 数据准备 --- # 假设你的文本文件名为 input.txt with open('input.txt', 'r', encoding='utf-8') as f: text = f.read() chars = sorted(list(set(text))) vocab_size = len(chars) print(f'Vocabulary size: {vocab_size}') char_to_idx = {ch: i for i, ch in enumerate(chars)} idx_to_char = {i: ch for i, ch in enumerate(chars)} data = torch.tensor([char_to_idx[ch] for ch in text], dtype=torch.long) def get_batch(split, batch_size, seq_length): split_ratio = 0.9 split_idx = int(len(data) * split_ratio) train_data = data[:split_idx] val_data = data[split_idx:] data_split = train_data if split == 'train' else val_data # 确保有足够的数据生成一个批次 if len(data_split) < seq_length + 1: raise ValueError(f"Data for split '{split}' is too short for sequence length {seq_length}.") starts = torch.randint(0, len(data_split) - seq_length, (batch_size,)) xs = torch.stack([data_split[start:start+seq_length] for start in starts]) ys = torch.stack([data_split[start+1:start+seq_length+1] for start in starts]) return xs, ys # --- 2. 模型定义 --- class CharGRU(nn.Module): def __init__(self, vocab_size, embed_dim, hidden_size, num_layers, dropout=0.2): super().__init__() self.embedding = nn.Embedding(vocab_size, embed_dim) self.gru = nn.GRU(embed_dim, hidden_size, num_layers, batch_first=True, dropout=dropout if num_layers>1 else 0) self.fc = nn.Linear(hidden_size, vocab_size) def forward(self, x, hidden=None): embedded = self.embedding(x) gru_out, hidden_out = self.gru(embedded, hidden) # 将GRU输出展平以通过全连接层 gru_out_flat = gru_out.contiguous().view(-1, gru_out.size(-1)) output = self.fc(gru_out_flat) # 重塑回 [batch, seq, vocab] 形状 output = output.view(gru_out.size(0), gru_out.size(1), -1) return output, hidden_out # --- 3. 文本生成函数 --- def generate_text(model, device, start_str, length=500, temperature=1.0): model.eval() chars = [ch for ch in start_str] input_seq = torch.tensor([[char_to_idx[ch] for ch in start_str]], dtype=torch.long).to(device) hidden = None with torch.no_grad(): for _ in range(length): logits, hidden = model(input_seq[:, -1:], hidden) # 只输入最后一个字符 logits = logits[:, -1, :] / temperature probs = torch.softmax(logits, dim=-1) next_idx = torch.multinomial(probs, num_samples=1).item() chars.append(idx_to_char[next_idx]) # 更新输入序列,只保留最后 seq_length 个字符以避免过长(可选) if input_seq.size(1) >= 100: # 保持一个合理的上下文窗口 input_seq = input_seq[:, 1:] input_seq = torch.cat([input_seq, torch.tensor([[next_idx]], device=device)], dim=1) return ''.join(chars) # --- 4. 训练配置 --- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f'Using device: {device}') embed_dim = 128 hidden_size = 512 # 稍微增大以增强表达能力 num_layers = 3 dropout = 0.2 learning_rate = 3e-4 batch_size = 64 seq_length = 256 # 增加序列长度以捕捉更长依赖 num_epochs = 10000 eval_interval = 500 print_interval = 100 model = CharGRU(vocab_size, embed_dim, hidden_size, num_layers, dropout).to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) # --- 5. 训练循环 --- model.train() for epoch in range(1, num_epochs + 1): xb, yb = get_batch('train', batch_size, seq_length) xb, yb = xb.to(device), yb.to(device) optimizer.zero_grad() logits, _ = model(xb) loss = criterion(logits.view(-1, vocab_size), yb.view(-1)) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() if epoch % print_interval == 0: print(f'Epoch {epoch:5d} | Loss: {loss.item():.4f}') if epoch % eval_interval == 0: model.eval() with torch.no_grad(): xb_val, yb_val = get_batch('val', batch_size, seq_length) xb_val, yb_val = xb_val.to(device), yb_val.to(device) logits_val, _ = model(xb_val) val_loss = criterion(logits_val.view(-1, vocab_size), yb_val.view(-1)) print(f'[Validation] Epoch {epoch:5d} | Val Loss: {val_loss.item():.4f}') # 生成示例文本 sample = generate_text(model, device, start_str='\nROMEO:', length=300, temperature=0.8) print('--- Generated Sample ---') print(sample) print('-' * 40) model.train() print('Training complete.') # --- 6. 最终生成示例 --- final_text = generate_text(model, device, start_str='\nThe future of AI is', length=1000, temperature=0.7) print('\n' + '='*50) print('FINAL GENERATED TEXT:') print('='*50) print(final_text)运行指南:
- 准备一个纯文本文件(如莎士比亚作品、小说、歌词等),命名为
input.txt,放在与脚本相同的目录下。 - 安装PyTorch(如前所述)。
- 运行脚本:
python train_rnn_text_generation.py。 - 观察控制台输出的损失值和定期生成的文本样本。训练可能需要数千到数万步才能看到连贯的文本。
- 你可以调整脚本顶部的超参数(
hidden_size,num_layers,learning_rate,temperature等)来观察对生成效果的影响。
训练这样一个模型,你最终会看到它从输出乱码开始,逐渐学会拼写单词,形成简单的短语,最终能生成具有基本语法结构和一定主题相关性的段落。这个过程本身,就是理解递归神经网络如何“学习”和“记忆”的最佳方式。
