动手学深度学习9.2. 长短期记忆网络（LSTM）-笔记&练习（PyTorch）

scdifsn

于 2024-10-08 15:27:14 发布

阅读量382

点赞数 11

分类专栏：《动手学深度学习》笔记及练习详解文章标签：深度学习 lstm 笔记 rnn

本文链接：https://blog.csdn.net/scdifsn/article/details/142760280

版权

《动手学深度学习》笔记及练习详解专栏收录该内容

54 篇文章 7 订阅

订阅专栏

本节课程地址：57 长短期记忆网络（LSTM）【动手学深度学习v2】_哔哩哔哩_bilibili

本节教材地址：9.2. 长短期记忆网络（LSTM） — 动手学深度学习 2.0.0 documentation (d2l.ai)

本节开源代码：...>d2l-zh>pytorch>chapter_multilayer-perceptrons>lstm.ipynb

长短期记忆网络（LSTM）

长期以来，隐变量模型存在着长期信息保存和短期输入缺失的问题。解决这一问题的最早方法之一是长短期存储器（long short-term memory，LSTM） (Hochreiter and Schmidhuber, 1997)。它有许多与门控循环单元（ 9.1节）一样的属性。有趣的是，长短期记忆网络的设计比门控循环单元稍微复杂一些，却比门控循环单元早诞生了近20年。

门控记忆元

可以说，长短期记忆网络的设计灵感来自于计算机的逻辑门。长短期记忆网络引入了记忆元（memory cell），或简称为单元（cell）。有些文献认为记忆元是隐状态的一种特殊类型，它们与隐状态具有相同的形状，其设计目的是用于记录附加的信息。为了控制记忆元，我们需要许多门。其中一个门用来从单元中输出条目，我们将其称为输出门（output gate）。另外一个门用来决定何时将数据读入单元，我们将其称为输入门（input gate）。我们还需要一种机制来重置单元的内容，由遗忘门（forget gate）来管理，这种设计的动机与门控循环单元相同，能够通过专用机制决定什么时候记忆或忽略隐状态中的输入。让我们看看这在实践中是如何运作的。

输入门、忘记门和输出门

就如在门控循环单元中一样，当前时间步的输入和前一个时间步的隐状态作为数据送入长短期记忆网络的门中，如图9.2.1 所示。它们由三个具有sigmoid激活函数的全连接层处理，以计算输入门、遗忘门和输出门的值。因此，这三个门的值都在 (0,1) 的范围内。

忘记门：降值朝0减少
输入门：决定是不是忽略掉输入数据
输出门：决定是不是使用隐状态

我们来细化一下长短期记忆网络的数学表达。假设有 $h$ 个隐藏单元，批量大小为 $n$ ，输入数为 $d$ 。因此，输入为 $\mathbf{X}_t \in \mathbb{R}^{n \times d}$ ，前一时间步的隐状态为 $\mathbf{H}_{t-1} \in \mathbb{R}^{n \times h}$ 。相应地，时间步 $t$ 的门被定义如下：输入门是 $\mathbf{I}_t \in \mathbb{R}^{n \times h}$ ，遗忘门是 $\mathbf{F}_t \in \mathbb{R}^{n \times h}$ ，输出门是 $\mathbf{O}_t \in \mathbb{R}^{n \times h}$ 。它们的计算方法如下：

$\begin{aligned} \mathbf{I}_t &= \sigma(\mathbf{X}_t \mathbf{W}_{xi} + \mathbf{H}_{t-1} \mathbf{W}_{hi} + \mathbf{b}_i),\\ \mathbf{F}_t &= \sigma(\mathbf{X}_t \mathbf{W}_{xf} + \mathbf{H}_{t-1} \mathbf{W}_{hf} + \mathbf{b}_f),\\ \mathbf{O}_t &= \sigma(\mathbf{X}_t \mathbf{W}_{xo} + \mathbf{H}_{t-1} \mathbf{W}_{ho} + \mathbf{b}_o), \end{aligned}$

其中 $\mathbf{W}_{xi}, \mathbf{W}_{xf}, \mathbf{W}_{xo} \in \mathbb{R}^{d \times h}$ 和 $\mathbf{W}_{hi}, \mathbf{W}_{hf}, \mathbf{W}_{ho} \in \mathbb{R}^{h \times h}$ 是权重参数， $\mathbf{b}_i, \mathbf{b}_f, \mathbf{b}_o \in \mathbb{R}^{1 \times h}$ 是偏置参数。

候选记忆元

由于还没有指定各种门的操作，所以先介绍候选记忆元（candidate memory cell） $\tilde{\mathbf{C}}_t \in \mathbb{R}^{n \times h}$ 。它的计算与上面描述的三个门的计算类似，但是使用 tanh 函数作为激活函数，函数的值范围为 (−1,1) 。下面导出在时间步 $t$ 处的方程：

$\tilde{\mathbf{C}}_t = \text{tanh}(\mathbf{X}_t \mathbf{W}_{xc} + \mathbf{H}_{t-1} \mathbf{W}_{hc} + \mathbf{b}_c),$

其中 $\mathbf{W}_{xc} \in \mathbb{R}^{d \times h}$ 和 $\mathbf{W}_{hc} \in \mathbb{R}^{h \times h}$ 是权重参数， $\mathbf{b}_c \in \mathbb{R}^{1 \times h}$ 是偏置参数。

候选记忆元的如图9.2.2 所示。

记忆元

在门控循环单元中，有一种机制来控制输入和遗忘（或跳过）。类似地，在长短期记忆网络中，也有两个门用于这样的目的：输入门 $\mathbf{I}_t$ 控制采用多少来自 $\tilde{\mathbf{C}}_t$ 的新数据，而遗忘门 $\mathbf{F}_t$ 控制保留多少过去的记忆元 $\mathbf{C}_{t-1} \in \mathbb{R}^{n \times h}$ 的内容。使用按元素乘法，得出：

$\mathbf{C}_t = \mathbf{F}_t \odot \mathbf{C}_{t-1} + \mathbf{I}_t \odot \tilde{\mathbf{C}}_t.$

如果遗忘门始终为 1 且输入门始终为 0 ，则过去的记忆元 $\mathbf{C}_{t-1}$ 将随时间被保存并传递到当前时间步。引入这种设计是为了缓解梯度消失问题，并更好地捕获序列中的长距离依赖关系。

这样我们就得到了计算记忆元的流程图，如图9.2.3。

隐状态

最后，我们需要定义如何计算隐状态 $\mathbf{H}_t \in \mathbb{R}^{n \times h}$ ，这就是输出门发挥作用的地方。在长短期记忆网络中，它仅仅是记忆元的 tanh 的门控版本。这就确保了 $\mathbf{H}_t$ 的值始终在区间 (−1,1) 内：

$\mathbf{H}_t = \mathbf{O}_t \odot \tanh(\mathbf{C}_t).$

只要输出门接近 1 ，我们就能够有效地将所有记忆信息传递给预测部分，而对于输出门接近 0 ，我们只保留记忆元内的所有信息，而不需要更新隐状态。

图9.2.4 提供了数据流的图形化演示。

从零开始实现

现在，我们从零开始实现长短期记忆网络。与 8.5节中的实验相同，我们首先加载时光机器数据集。

import torch
from torch import nn
from d2l import torch as d2l

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

[初始化模型参数]

接下来，我们需要定义和初始化模型参数。如前所述，超参数num_hiddens定义隐藏单元的数量。我们按照标准差 0.01 的高斯分布初始化权重，并将偏置项设为 0 。

def get_lstm_params(vocab_size, num_hiddens, device):
    num_inputs = num_outputs = vocab_size

    def normal(shape):
        return torch.randn(size=shape, device=device)*0.01

    def three():
        return (normal((num_inputs, num_hiddens)),
                normal((num_hiddens, num_hiddens)),
                torch.zeros(num_hiddens, device=device))

    W_xi, W_hi, b_i = three()  # 输入门参数
    W_xf, W_hf, b_f = three()  # 遗忘门参数
    W_xo, W_ho, b_o = three()  # 输出门参数
    W_xc, W_hc, b_c = three()  # 候选记忆元参数
    # 输出层参数
    W_hq = normal((num_hiddens, num_outputs))
    b_q = torch.zeros(num_outputs, device=device)
    # 附加梯度
    params = [W_xi, W_hi, b_i, W_xf, W_hf, b_f, W_xo, W_ho, b_o, W_xc, W_hc,
              b_c, W_hq, b_q]
    for param in params:
        param.requires_grad_(True)
    return params

定义模型

在[初始化函数]中，长短期记忆网络的隐状态需要返回一个额外的记忆元，单元的值为0，形状为（批量大小，隐藏单元数）。因此，我们得到以下的状态初始化。

def init_lstm_state(batch_size, num_hiddens, device):
    return (torch.zeros((batch_size, num_hiddens), device=device),
            torch.zeros((batch_size, num_hiddens), device=device))

[实际模型]的定义与我们前面讨论的一样：提供三个门和一个额外的记忆元。请注意，只有隐状态才会传递到输出层，而记忆元 $\mathbf{C}_t$ 不直接参与输出计算。

def lstm(inputs, state, params):
    [W_xi, W_hi, b_i, W_xf, W_hf, b_f, W_xo, W_ho, b_o, W_xc, W_hc, b_c,
     W_hq, b_q] = params
    (H, C) = state
    outputs = []
    for X in inputs:
        I = torch.sigmoid((X @ W_xi) + (H @ W_hi) + b_i)
        F = torch.sigmoid((X @ W_xf) + (H @ W_hf) + b_f)
        O = torch.sigmoid((X @ W_xo) + (H @ W_ho) + b_o)
        C_tilda = torch.tanh((X @ W_xc) + (H @ W_hc) + b_c)
        C = F * C + I * C_tilda
        H = O * torch.tanh(C)
        Y = (H @ W_hq) + b_q
        outputs.append(Y)
    return torch.cat(outputs, dim=0), (H, C)

[训练]和预测

让我们通过实例化 8.5节中引入的RNNModelScratch类来训练一个长短期记忆网络，就如我们在 9.1节中所做的一样。

vocab_size, num_hiddens, device = len(vocab), 256, d2l.try_gpu()
num_epochs, lr = 500, 1
model = d2l.RNNModelScratch(len(vocab), num_hiddens, device, get_lstm_params,
                            init_lstm_state, lstm)
d2l.train_ch8(model, train_iter, vocab, lr, num_epochs, device)

输出结果：
perplexity 1.1, 18890.1 tokens/sec on cpu
time traveller after the pausere ulh s i wonyrever a tountho mom
traveller for throught yoand monemittely is what insearn an

[简洁实现]

使用高级API，我们可以直接实例化LSTM模型。高级API封装了前文介绍的所有配置细节。这段代码的运行速度要快得多，因为它使用的是编译好的运算符而不是Python来处理之前阐述的许多细节。

num_inputs = vocab_size
lstm_layer = nn.LSTM(num_inputs, num_hiddens)
model = d2l.RNNModel(lstm_layer, len(vocab))
model = model.to(device)
d2l.train_ch8(model, train_iter, vocab, lr, num_epochs, device)

输出结果：
perplexity 1.0, 116254.4 tokens/sec on cpu
time travelleryou can show black is white by argument said filby
travelleryou can show black is white by argument said filby

长短期记忆网络是典型的具有重要状态控制的隐变量自回归模型。多年来已经提出了其许多变体，例如，多层、残差连接、不同类型的正则化。然而，由于序列的长距离依赖性，训练长短期记忆网络和其他序列模型（例如门控循环单元）的成本是相当高的。在后面的内容中，我们将讲述更高级的替代模型，如Transformer。

小结

长短期记忆网络有三种类型的门：输入门、遗忘门和输出门。
长短期记忆网络的隐藏层输出包括“隐状态”和“记忆元”。只有隐状态会传递到输出层，而记忆元完全属于内部信息。
长短期记忆网络可以缓解梯度消失和梯度爆炸。

练习

调整和分析超参数对运行时间、困惑度和输出顺序的影响。

解：
类似9.1节的练习第2题，超参数包括：

批量大小 batch_size，由下面代码试验结果0和1比较可以看出：batch_size增加会增加运行时间，并大大增加困惑度。
批量步长 num_steps，由下面代码试验结果0和2比较可以看出：num_steps降低会略增加运行时间，但困惑度降为最低。
隐藏层数 num_hiddens，由下面代码试验结果0和3比较可以看出：num_hiddens增加会大大增加运行时间，困惑度不变，但从输出来看，增加num_hiddens后的输出字符组成的word比较自然。
迭代次数 num_epochs，这个可以从运行结果图中直接看出：epochs越大，运行时间越长；如果epochs过低，则模型无法完成收敛，困惑度不能有效降低，因此，需要保证足够的epochs。下面代码不做epochs的调整。
学习率 lr，由下面代码试验结果0和4比较可以看出：learning rate降低会导致模型无法完成收敛，困惑度很高。

代码如下：

import time

# 单一变量依次调整 batch_size, num_steps, num_hiddens, lr
hyper_0 = [32, 35, 256, 1] # 初始值
hyper_1 = [64, 35, 256, 1]
hyper_2 = [32, 10, 256, 1]
hyper_3 = [32, 35, 512, 1]
hyper_4 = [32, 35, 256, 0.1]
hyper_params = [hyper_0, hyper_1, hyper_2, hyper_3, hyper_4]

def train_ch9_2(net, train_iter, vocab, lr, num_epochs, device,
              use_random_iter=False):
    loss = nn.CrossEntropyLoss()
    # 初始化
    if isinstance(net, nn.Module):
        updater = torch.optim.SGD(net.parameters(), lr)
    else:
        updater = lambda batch_size: d2l.sgd(net.params, lr, batch_size)
    predict = lambda prefix: d2l.predict_ch8(prefix, 50, net, vocab, device)
    # 训练和预测
    for epoch in range(num_epochs):
        ppl, _ = d2l.train_epoch_ch8(
            net, train_iter, loss, updater, device, use_random_iter)
    print(predict('time traveller'))
    print(predict('traveller'))
    return ppl 

ppls = []
run_times = []

for i, hyper in enumerate(hyper_params):
    num_epochs = 500
    batch_size, num_steps, num_hiddens, lr = hyper
    train_iter, vocab = d2l.load_data_time_machine(batch_size, num_steps)
    vocab_size, device = len(vocab), d2l.try_gpu()
    model = d2l.RNNModelScratch(len(vocab), num_hiddens, device, get_lstm_params,
                            init_lstm_state, lstm)
    print(f'hyper_{i}:')
    print(f'When batch_size = {hyper[0]}, num_steps = {hyper[1]}, num_hiddens = {hyper[2]}, lr = {hyper[3]}:')
    start = time.time()
    ppl = train_ch9_2(model, train_iter, vocab, lr, num_epochs, device)
    end = time.time()
    run_time = end - start
    print(f'困惑度 {ppl:.1f}, {run_time:.1f} 秒')

    ppls.append(ppl)
    run_times.append(run_time)

输出结果：
hyper_0:
When batch_size = 32, num_steps = 35, num_hiddens = 256, lr = 1:
time traveller whthef about this wilb ip time are arausteryshis
traveller chard and and anous to kan ofer than a small cloc
困惑度 1.1, 245.1 秒
hyper_1:
When batch_size = 64, num_steps = 35, num_hiddens = 256, lr = 1:
time traveller thing the time traveller thing the time traveller
traveller thing the time traveller thing the time traveller
困惑度 5.8, 396.6 秒
hyper_2:
When batch_size = 32, num_steps = 10, num_hiddens = 256, lr = 1:
time travelleryou can show black is white by argument said filby
travelleryou can show black is white by argument said filby
困惑度 1.0, 293.0 秒
hyper_3:
When batch_size = 32, num_steps = 35, num_hiddens = 512, lr = 1:
time traveller for so it will be convenient to speak of himwas e
travelleryou can show black is white by argument said filby
困惑度 1.1, 841.9 秒
hyper_4:
When batch_size = 32, num_steps = 35, num_hiddens = 256, lr = 0.1:
time traveller at at at at at at at at at at at at at at at at a
traveller at at at at at at at at at at at at at at at at a
困惑度 14.3, 256.4 秒

import matplotlib.pyplot as plt

fig, ax1 = plt.subplots(figsize=(8, 6))
ax2 = ax1.twinx()
ax1.bar([0, 1, 2, 3, 4], ppls, alpha=0.8)
ax1.set_ylim()
ax1.set_xlabel('No. hyper_params')
ax1.set_ylabel('Perplexity')
ax2.plot([0, 1, 2, 3, 4], run_times, linewidth='2', color = 'darkred')
ax2.set_ylabel('Run_time')
ax2.set_ylim()
plt.plot(0, 0, linewidth='2', color = 'darkred', label="Run_time")
plt.bar(0, 0, alpha=0.8, label="Perplexity")
plt.legend(loc='upper left', frameon=False)
plt.show()

2. 如何更改模型以生成适当的单词，而不是字符序列？

解：
词元化处理时改为按'word'分词，而不是按'char'分词；模型预测和训练也做相应调整。
代码如下：

# 改为'word'分词
def load_corpus_time_machine_word(max_tokens=-1):
    lines = d2l.read_time_machine()
    tokens = d2l.tokenize(lines, 'word')
    vocab = d2l.Vocab(tokens)
    corpus = [vocab[token] for line in tokens for token in line]
    if max_tokens > 0: 
        corpus = corpus[:max_tokens]
    return corpus, vocab

class SeqDataLoader_word:
    def __init__(self, batch_size, num_steps, use_random_iter, max_tokens):
        if use_random_iter:
            self.data_iter_fn = d2l.seq_data_iter_random
        else:
            self.data_iter_fn = d2l.seq_data_iter_sequential
        self.corpus, self.vocab = load_corpus_time_machine_word(max_tokens)
        self.batch_size, self.num_steps = batch_size, num_steps

    def __iter__(self):
        return self.data_iter_fn(self.corpus, self.batch_size, self.num_steps)

def load_data_time_machine_word(batch_size, num_steps, 
                           use_random_iter=False, max_tokens=10000):
    data_iter = SeqDataLoader_word(
        batch_size, num_steps, use_random_iter, max_tokens)
    return data_iter, data_iter.vocab
def predict_word(prefix, num_preds, net, vocab, device):
    state = net.begin_state(batch_size=1, device=device)
    # 对prefix进行word分割
    prefix = prefix.strip().lower().split()
    outputs = [vocab[prefix[0]]]
    get_input = lambda: d2l.reshape(d2l.tensor(
        [outputs[-1]], device=device), (1, 1))
    for y in prefix[1:]:  # Warm-up period
        _, state = net(get_input(), state)
        outputs.append(vocab[y])
    for _ in range(num_preds):  # Predict `num_preds` steps
        y, state = net(get_input(), state)
        outputs.append(int(y.argmax(dim=1).reshape(1)))
    # 这里改为空格join，不然输出的word中间没有空格隔开
    return ' '.join([vocab.idx_to_token[i] for i in outputs])
def train_ch9_2_word(net, train_iter, vocab, lr, num_epochs, device,
              use_random_iter=False):
    loss = nn.CrossEntropyLoss()
    animator = d2l.Animator(xlabel='epoch', ylabel='perplexity',
                            legend=['train'], xlim=[10, num_epochs])
    # 初始化
    if isinstance(net, nn.Module):
        updater = torch.optim.SGD(net.parameters(), lr)
    else:
        updater = lambda batch_size: d2l.sgd(net.params, lr, batch_size)
    predict = lambda prefix: predict_word(prefix, 50, net, vocab, device)
    # 训练和预测
    for epoch in range(num_epochs):
        ppl, speed = d2l.train_epoch_ch8(
            net, train_iter, loss, updater, device, use_random_iter)
        if (epoch + 1) % 10 == 0:
            print(predict('time traveller'))
            animator.add(epoch + 1, [ppl])
    print(f'困惑度 {ppl:.1f}, {speed:.1f} 词元/秒 {str(device)}')
    print(predict('time traveller'))
    print(predict('traveller'))
batch_size, num_steps = 32, 10
train_iter, vocab = load_data_time_machine_word(batch_size, num_steps)
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 = d2l.RNNModel(lstm_layer, len(vocab))
model = model.to(device)
train_ch9_2_word(model, train_iter, vocab, lr, num_epochs, device)

输出结果：
困惑度 1.0, 8117.7 词元/秒 cpu
time traveller for his eyelashes the silent man seemed even more clumsy than usual and drank champagne with regularity and determination out of sheer nervousness at last the time traveller pushed his plate away and looked round us i suppose i must apologize he said i was simply starving i ve had
traveller and his hard to readjust it as i did so the shafts of the sun smote through the thunderstorm the grey downpour was swept aside and vanished like the trailing garments of a ghost above me in the intense blue of the summer sky some faint brown shreds of cloud

3. 在给定隐藏层维度的情况下，比较门控循环单元、长短期记忆网络和常规循环神经网络的计算成本。要特别注意训练和推断成本。

解：
固定所有超参数，包括隐藏层维度，比较对同一数据集分别使用RNN、GRU、LSTM模型的计算成本。
试验结果表明：

RNN的计算用时最短，因为RNN模型结构最简单，参数也最少，但是困惑度也最高，GRU和LSTM的困惑度均降到最低为1.0。
有趣的是，在CPU运行时，虽然LSTM比GRU的参数量更多，但是在给定的超参数和数据库条件下，LSTM的计算用时比GRU更少。
代码如下：

import time

batch_size, num_steps = 32, 35
train_iter, vocab = d2l.load_data_time_machine(batch_size, num_steps)
vocab_size, num_hiddens, device = len(vocab), 256, d2l.try_gpu()
num_epochs, lr = 500, 1
num_inputs = vocab_size
# RNN
rnn_layer = nn.RNN(num_inputs, num_hiddens)
net = d2l.RNNModel(rnn_layer, vocab_size=len(vocab))
net = net.to(device)
start = time.time()
d2l.train_ch8(net, train_iter, vocab, lr, num_epochs, device)
end = time.time()

print(f'RNN运行时长：{end-start:.1f}')

输出结果：
perplexity 1.3, 79589.5 tokens/sec on cpu
time traveller ffocletsay it flathat saightwit suthe weysure tha
traveller hepecenath se then th aratile the end of oprac us
RNN运行时长：65.1

# GRU
gru_layer = nn.GRU(num_inputs, num_hiddens)
net = d2l.RNNModel(gru_layer, vocab_size=len(vocab))
net = net.to(device)
start = time.time()
d2l.train_ch8(net, train_iter, vocab, lr, num_epochs, device)
end = time.time()
print(f'GRU运行时长：{end-start:.1f}')

输出结果：
perplexity 1.0, 40279.4 tokens/sec on cpu
time travelleryou can show black is white by argument said filby
travelleryou can show black is white by argument said filby
GRU运行时长：123.4

# LSTM
lstm_layer = nn.LSTM(num_inputs, num_hiddens)
net = d2l.RNNModel(lstm_layer, vocab_size=len(vocab))
net = net.to(device)
start = time.time()
d2l.train_ch8(net, train_iter, vocab, lr, num_epochs, device)
end = time.time()
print(f'LSTM运行时长：{end-start:.1f}')

输出结果：
perplexity 1.0, 52762.0 tokens/sec on cpu
time travelleryou can show black is white by argument said filby
travelleryou can show black is white by argument said filby
LSTM运行时长：91.1

3. 既然候选记忆元通过使用 tanh 函数来确保值范围在 (−1,1) 之间，那么为什么隐状态需要再次使用 tanh 函数来确保输出值范围在 (−1,1) 之间呢？
解：
虽然候选记忆元通过使用 tanh 函数来确保值范围在 (−1,1) 之间，但计算记忆元的时候是 $\mathbf{C}_t = \mathbf{F}_t \odot \mathbf{C}_{t-1} + \mathbf{I}_t \odot \tilde{\mathbf{C}}_t$ ，使得记忆元的取值范围扩大到$ (−2,2) 之间，因此，计算隐状态时需要再次使用 tanh 函数来确保输出值范围在 (−1,1) 之间。

4. 实现一个能够基于时间序列进行预测而不是基于字符序列进行预测的长短期记忆网络模型。
解：
用LSTM实现8.1节的自回归模型，基于时间序列进行预测。代码如下：

# 生成序列数据和训练模型
T = 1000 
time = torch.arange(1, T + 1, dtype=torch.float32)
x = torch.sin(0.01 * time) + torch.normal(0, 0.2, (T,))

tau = 4
features = torch.zeros((T - tau, tau))
for i in range(tau):
    features[:, i] = x[i: T - tau + i]
labels = x[tau:].reshape((-1, 1))

batch_size, n_train = 16, 600

train_iter = d2l.load_array((features[:n_train], labels[:n_train]),
                            batch_size, is_train=True)
# 构建LSTM和训练函数
class RNNModel(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size):
        super(RNNModel, self).__init__()
        # 构建隐藏层+输出层
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers)
        self.linear = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        # x的形状应为(batch_size, sequence_length, input_size)
        # batch_size = 16, sequence_length = 1, input_size = 4
        Y, _ = self.lstm(x.unsqueeze(1))
        # x的形状应为(batch_size, sequence_length, hidden_size)
        output = self.linear(Y[:, -1, :])
        return output

def train(net, train_iter, epochs, lr):
    trainer = torch.optim.Adam(net.parameters(), lr)
    loss = nn.MSELoss(reduction='none')
    for epoch in range(epochs):
        for X, y in train_iter:
            trainer.zero_grad()
            l = loss(net(X), y)
            l.sum().backward()
            trainer.step()
        print(f'epoch {epoch + 1}, '
              f'loss: {d2l.evaluate_loss(net, train_iter, loss):f}')
# 训练
input_size, hidden_size, num_layers, output_size = 4, 10, 1, 1
net_seq = RNNModel(input_size, hidden_size, num_layers, output_size)

train(net_seq, train_iter, 5, 0.01)

输出结果：
epoch 1, loss: 0.055929
epoch 2, loss: 0.048768
epoch 3, loss: 0.046579
epoch 4, loss: 0.048393
epoch 5, loss: 0.043355

# 预测
onestep_preds = net_seq(features)
d2l.plot([time, time[tau:]],
         [x.detach().numpy(), onestep_preds.detach().numpy()], 'time',
         'x', legend=['data', '1-step preds'], xlim=[1, 1000],
         figsize=(6, 3))