代码实现LSTM 长短期记忆网络 | #51CTO博主之星评选#

Posted 2022-03-29 LolitaAnn

之前我还会写手动实现，就是实现以下计算过程，然而实际上其实就是用代码堆出来计算公式，也没什么意思，:weary:以后就不搞了，直接写怎么用pytorch实现。

---

import torch
from torch import nn
from d2l import torch as d2l
from torch.nn import functional as F

导包啊导包，这个不用解释了吧。

train_iter, vocab = d2l.load_data_time_machine(batch_size, num_steps)
batch_size, num_steps = 32, 35

这里直接使用d2l加载数据集。

设定批量大小batch_size和时间步的长度num_steps，时间步的长度就是每次LSTM处理的一个序列的长度。

class LSTMModel(nn.Module):
    def __init__(self, lstm_layer, vocab_size, **kwargs):
        super(LSTMModel, self).__init__(**kwargs)
        self.lstm = lstm_layer
        self.vocab_size = vocab_size
        self.num_hiddens = self.lstm.hidden_size
        if not self.lstm.bidirectional:
            self.num_directions = 1
            self.linear = nn.Linear(self.num_hiddens, self.vocab_size)
        else:
            self.num_directions = 2
            self.linear = nn.Linear(self.num_hiddens * 2, self.vocab_size)

    def forward(self, inputs, state):
        X = F.one_hot(inputs.T.long(), self.vocab_size)
        X = X.to(torch.float32)
        Y, state = self.lstm(X, state)
        output = self.linear(Y.reshape((-1, Y.shape[-1])))
        return output, state

    def begin_state(self, device, batch_size=1):
        return (torch.zeros((
            self.num_directions * self.lstm.num_layers,
            batch_size, self.num_hiddens), device=device),
                torch.zeros((
                    self.num_directions * self.lstm.num_layers,
                    batch_size, self.num_hiddens), device=device))

• __init__初始化这个类，这个类是继承了nn.Module的。
  • self.lstm设定计算层是LSTM层
  • self.vocab_size设定字典的大小，这里大小是28，因为为了方便演示，我们这里使用的是字母进行分词，所以其中只有a\\~z26个字母外加`和<unk>`（空格和unknown）。
  • self.num_hiddens设置隐藏层的大小。
  • if-else语句是设定LSTM单双向的，毕竟还有双向LSTM这种东西的存在。
• forward定义前向传播网络。
  也就是描述计算过程。
  • 这里首先是将输入转化为对应的one-hot向量，再将其类型转化为float。
  • Y和state是计算隐状态的，这里Y是输出全部的隐状态，state是输出最后一个时间步的隐状态。注意 在这里Y不是 输出。
  • output是用于存储输出的。
• begin_state是进行初始化。
  这里是return了好长一个句子。可以拆解开看一下子。
  
  这里是初始化为0张量，初始化位置device=device由你传入的位置决定是CPU还是GPU。这里和普通RNN的区别在于普通RNN和GRU不同，二者只需要返回一个张量即可，但LSTM这里是一个元组里两个张量。

这段代码看似是写了个LSTM的类，其实是换汤不换药的，就是之前手动简洁实现RNN那个文章里的RNN类改了一下子。不论是RNN还是GRU、LSTM，都是在那个类的基础上改的。那个RNN的类写的更齐全，详细的可以看:point_right:简洁实现RNN循环神经网络

vocab_size, num_hiddens, device = len(vocab), 256, d2l.try_gpu()
num_epochs, lr = 500, 1
num_inputs = vocab_size
lstm_layer = nn.LSTM(num_inputs, num_hiddens)
model = LSTMModel(lstm_layer, len(vocab))
model = model.to(device)
d2l.train_ch8(model, train_iter, vocab, lr, num_epochs, device)

然后就是我们的训练和预测过程。

• 第一句和第二句分别是设置词典长度、隐藏层大小、运行在CPU还是GPU上、训练epoch数量、学习率。
• 设置lstm使用pytorch自带的nn.LSTM。
• 模型使用我们的那个类，并且将其放到对应的设备上运行。
• 最后一句就是预测训练的过程。

下图是100个epoch之后的结果，还没降到底，我这组训练数据要跑400多才能差不多平稳。结果会输出困惑度以及前缀为“time traveler”和“traveler”的预测结果。