PyTorch--语言建模的循环神经网络（RNN）的实现

前言

嗨嗨嗨，这是每日代码小记，今天的代码是一个用于语言建模的循环神经网络（RNN）的PyTorch实现，灵感来源于PyTorch官方示例仓库中的"word_language_model"。以下是代码的关键部分的详细解析：

完整代码

# Some part of the code was referenced from below.
# https://github.com/pytorch/examples/tree/master/word_language_model 
import torch
import torch.nn as nn
import numpy as np
from torch.nn.utils import clip_grad_norm_
from data_utils import Dictionary, Corpus


# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Hyper-parameters
embed_size = 128
hidden_size = 1024
num_layers = 1
num_epochs = 5
num_samples = 1000     # number of words to be sampled
batch_size = 20
seq_length = 30
learning_rate = 0.002

# Load "Penn Treebank" dataset
corpus = Corpus()
ids = corpus.get_data('data/train.txt', batch_size)
vocab_size = len(corpus.dictionary)
num_batches = ids.size(1) // seq_length


# RNN based language model
class RNNLM(nn.Module):
    def __init__(self, vocab_size, embed_size, hidden_size, num_layers):
        super(RNNLM, self).__init__()
        self.embed = nn.Embedding(vocab_size, embed_size)
        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True)
        self.linear = nn.Linear(hidden_size, vocab_size)
        
    def forward(self, x, h):
        # Embed word ids to vectors
        x = self.embed(x)
        
        # Forward propagate LSTM
        out, (h, c) = self.lstm(x, h)
        
        # Reshape output to (batch_size*sequence_length, hidden_size)
        out = out.reshape(out.size(0)*out.size(1), out.size(2))
        
        # Decode hidden states of all time steps
        out = self.linear(out)
        return out, (h, c)

model = RNNLM(vocab_size, embed_size, hidden_size, num_layers).to(device)

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# Truncated backpropagation
def detach(states):
    return [state.detach() for state in states] 

# Train the model
for epoch in range(num_epochs):
    # Set initial hidden and cell states
    states = (torch.zeros(num_layers, batch_size, hidden_size).to(device),
              torch.zeros(num_layers, batch_size, hidden_size).to(device))
    
    for i in range(0, ids.size(1) - seq_length, seq_length):
        # Get mini-batch inputs and targets
        inputs = ids[:, i:i+seq_length].to(device)
        targets = ids[:, (i+1):(i+1)+seq_length].to(device)
        
        # Forward pass
        states = detach(states)
        outputs, states = model(inputs, states)
        loss = criterion(outputs, targets.reshape(-1))
        
        # Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        clip_grad_norm_(model.parameters(), 0.5)
        optimizer.step()

        step = (i+1) // seq_length
        if step % 100 == 0:
            print ('Epoch [{}/{}], Step[{}/{}], Loss: {:.4f}, Perplexity: {:5.2f}'
                   .format(epoch+1, num_epochs, step, num_batches, loss.item(), np.exp(loss.item())))

# Test the model
with torch.no_grad():
    with open('sample.txt', 'w') as f:
        # Set intial hidden ane cell states
        state = (torch.zeros(num_layers, 1, hidden_size).to(device),
                 torch.zeros(num_layers, 1, hidden_size).to(device))

        # Select one word id randomly
        prob = torch.ones(vocab_size)
        input = torch.multinomial(prob, num_samples=1).unsqueeze(1).to(device)

        for i in range(num_samples):
            # Forward propagate RNN 
            output, state = model(input, state)

            # Sample a word id
            prob = output.exp()
            word_id = torch.multinomial(prob, num_samples=1).item()

            # Fill input with sampled word id for the next time step
            input.fill_(word_id)

            # File write
            word = corpus.dictionary.idx2word[word_id]
            word = '\n' if word == '<eos>' else word + ' '
            f.write(word)

            if (i+1) % 100 == 0:
                print('Sampled [{}/{}] words and save to {}'.format(i+1, num_samples, 'sample.txt'))

# Save the model checkpoints
torch.save(model.state_dict(), 'model.ckpt')

代码分析

导入必要的库

import torch
import torch.nn as nn
import numpy as np
from torch.nn.utils import clip_grad_norm_
from data_utils import Dictionary, Corpus

导入PyTorch及其相关模块，以及用于数据处理的自定义模块。

设备配置

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

根据是否有可用的GPU，设置计算设备。

超参数设置

embed_size = 128
hidden_size = 1024
num_layers = 1
num_epochs = 5
...

设置模型和训练过程的超参数。

数据加载

corpus = Corpus()
ids = corpus.get_data('data/train.txt', batch_size)
vocab_size = len(corpus.dictionary)
...

加载并处理"Penn Treebank"数据集，创建词汇表和数据批次。

定义RNN模型

class RNNLM(nn.Module):
    ...
    def forward(self, x, h):
        ...

定义一个基于LSTM的RNN语言模型，包含嵌入层、LSTM层和线性层。

实例化模型并移动到设备

model = RNNLM(vocab_size, embed_size, hidden_size, num_layers).to(device)

创建模型实例并将模型移动到配置的设备上。

损失函数和优化器

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

定义交叉熵损失函数和Adam优化器。

截断反向传播

def detach(states):
    return [state.detach() for state in states]

定义一个函数，用于在反向传播中截断LSTM的状态。

训练模型

for epoch in range(num_epochs):
    ...
    for i in range(0, ids.size(1) - seq_length, seq_length):
        ...

执行模型的训练循环，包括数据准备、前向传播、损失计算、反向传播和参数更新。

测试模型

with torch.no_grad():
    ...

在测试模式下，使用模型生成新的文本样本并保存到文件。

保存模型

torch.save(model.state_dict(), 'model.ckpt')

保存训练好的模型参数。

这段代码实现了一个完整的训练和测试流程，适合用于语言建模任务。通过调整超参数，您可以探索不同模型配置对性能的影响。

亮点

这段代码在处理变长序列时主要使用了以下技术：

1. 截断反向传播（Truncated Backpropagation）：

   • 通过将长序列分割成固定长度的批次进行训练，以适应RNN的序列处理能力。这种方法允许模型一次只处理序列的一部分，而不是整个序列。

2. 可变长度序列的批处理：

   • 在每个epoch中，代码通过ids.size(1) - seq_length来确定序列的结束位置，并以seq_length为步长遍历整个序列。这允许模型处理不同长度的序列。

3. 隐藏状态的分离（Detaching States）：

   • 使用detach函数来分离隐藏状态，确保在反向传播过程中不会对之前的序列步骤产生梯度。这是通过调用.detach()方法实现的，它创建了隐藏状态的副本，这些副本不会参与梯度计算。

4. 序列长度的动态调整：

   • 代码中的num_batches = ids.size(1) // seq_length计算了在给定序列长度下，数据集中可以形成多少个批次。这允许模型处理不同长度的序列，同时保持批次的一致性。

5. 重塑输出（Reshaping Outputs）：

   • 在前向传播过程中，模型的输出被重塑为(batch_size * sequence_length, hidden_size)的形状，以便在解码时使用所有时间步的隐藏状态。

6. 逐元素损失计算：

   • 使用criterion(outputs, targets.reshape(-1))计算损失时，目标张量被重塑为一维，以匹配输出张量的形状。这允许模型逐元素地计算损失，而不是整个序列。

7. 随机采样生成文本：

   • 在生成新文本时，模型使用torch.multinomial对每个时间步的输出进行采样，从而生成新的词ID。这种方法适用于生成任意长度的序列。

这些技术共同使得代码能够有效地处理变长序列，并在训练和生成文本时保持灵活性。

模型小介绍

使用的模型是一个基于长短期记忆网络（LSTM）的循环神经网络（RNN），专门用于语言建模任务。以下是模型的关键组件和特点的介绍：

1. 模型类定义 (RNNLM)

• RNNLM类继承自nn.Module，是PyTorch中定义自定义神经网络的基类。

2. 嵌入层 (self.embed)

• 使用nn.Embedding实现，将词汇表中的单词ID映射到连续的向量空间中。这一层是词汇表大小乘以嵌入大小的矩阵。

3. LSTM层 (self.lstm)

• 使用nn.LSTM实现，是模型的核心，负责处理序列数据并捕获时间序列中的动态特征。LSTM层可以包含多个记忆单元，每个单元可以学习序列中不同时间步长的信息。

4. 线性层 (self.linear)

• 使用nn.Linear实现，将LSTM层的输出转换为词汇表大小的向量，以便进行下一步的预测。

5. 前向传播 (forward 方法)

• 在forward方法中，首先将输入的单词ID通过嵌入层转换为向量形式。
• 然后，这些向量被送入LSTM层进行序列处理，得到每个时间步的隐藏状态。
• 最后，通过线性层将隐藏状态映射到词汇表大小的输出上，这些输出可以通过softmax函数转换为概率分布，用于预测下一个单词。

6. 隐藏状态和细胞状态 (states)

• LSTM层的输出包括隐藏状态和细胞状态，这些状态可以跨时间步传递，帮助模型记住长距离依赖关系。

7. 截断反向传播

• 为了处理长序列，使用截断反向传播来限制计算资源的使用和梯度的流动。这意味着在每个批次的序列上只进行固定步数的反向传播。

8. 梯度裁剪

• 使用clip_grad_norm_函数来限制梯度的大小，防止梯度爆炸问题，这是一种常见的稳定训练过程的技术。

9. 损失函数

• 使用nn.CrossEntropyLoss作为损失函数，它计算模型输出和目标之间的交叉熵，适用于分类问题。

10. 优化器

- 使用`torch.optim.Adam`作为优化器，它是一种自适应学习率的优化算法，通常在训练深度学习模型时表现良好。

11. 文本生成

- 模型训练完成后，可以使用它来生成新的文本。通过从词汇表中采样单词并使用模型预测下一个单词，可以生成连贯的文本序列。

模型的构成在代码中由RNNLM类实现，具体细节如下：

模型类：RNNLM

class RNNLM(nn.Module):
    def __init__(self, vocab_size, embed_size, hidden_size, num_layers):
        super(RNNLM, self).__init__()
        self.embed = nn.Embedding(vocab_size, embed_size)
        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True)
        self.linear = nn.Linear(hidden_size, vocab_size)

组成部分：

1. 嵌入层 (self.embed):

   • nn.Embedding(vocab_size, embed_size)
   • 将词汇表中的每个单词ID映射到一个固定大小的嵌入向量。

2. LSTM层 (self.lstm):

   • nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True)
   • 一个长短期记忆网络层，用于处理序列数据。
   • embed_size: 输入特征的维度，与嵌入层的输出维度相同。
   • hidden_size: 隐藏层的维度。
   • num_layers: 堆叠的LSTM层数。
   • batch_first=True: 表示输入和输出张量的第一个维度是批次大小。

3. 线性层 (self.linear):

   • nn.Linear(hidden_size, vocab_size)
   • 将LSTM层的输出映射到词汇表大小的向量空间，用于预测下一个单词的概率分布。

前向传播：forward 方法

def forward(self, x, h):
    x = self.embed(x)
    out, (h, c) = self.lstm(x, h)
    out = out.reshape(out.size(0) * out.size(1), out.size(2))
    out = self.linear(out)
    return out, (h, c)

方法的工作流程：

1. 嵌入单词ID: 输入的单词ID通过嵌入层转换为嵌入向量。
2. LSTM前向传播: 嵌入向量通过LSTM层，同时传入隐藏状态h和细胞状态c，并更新这些状态。
3. 重塑输出: LSTM的输出被重塑为(batch_size * sequence_length, hidden_size)的形状。
4. 线性变换: 重塑后的输出通过线性层进行变换，得到每个时间步的预测分布。
5. 返回值: 返回模型的输出和更新后的状态。

这个RNNLM模型是一个典型的序列到序列（Seq2Seq）模型，适用于语言建模和其他需要处理序列数据的任务。通过LSTM层，模型能够捕捉序列中的长期依赖关系，并通过线性层生成下一个单词的概率分布。