基于LSTM算法实现交通流量预测(Pytorch版)

算法介绍

LSTM（Long Short-Term Memory）算法是一种特殊设计的循环神经网络（RNN, Recurrent Neural Network），专为有效地处理和建模序列数据中的长期依赖关系而开发。由于传统RNN在处理长序列时容易遇到梯度消失和梯度爆炸问题，导致模型难以捕捉到远距离输入之间的关联，LSTM通过引入独特的细胞状态（cell state）和多层门控机制解决了这些问题，从而在各种序列学习任务中展现出强大的性能。

LSTM模型的结构如下:

LSTM的关键组件如下：

遗忘门

LSTM（Long Short-Term Memory）网络的遗忘门（Forget Gate）是其门控机制的关键部分之一，负责决定在给定时间步 t t t 时，上一时刻的细胞状态 C t − 1 C_{t-1} Ct−1 中哪些信息应该被保留，哪些应该被遗忘。遗忘门的工作流程如下：

遗忘门计算

1. 输入合并 ：

   遗忘门接收前一时刻隐藏状态 h t − 1 h_{t-1} ht−1 和当前时间步输入 x t x_t xt。这两个向量被拼接（concatenate）或通过某种形式的线性变换（如全连接层）组合在一起，形成一个综合的输入向量，表示当前时间步的全部可用信息。

2. 门控值计算 ：

   综合输入向量经过一个带有sigmoid激活函数的全连接层（或称线性层），计算出遗忘门的输出值 f t f_t ft。sigmoid函数将这个综合输入向量映射到一个介于0和1之间的值，其中：

   • 值接近0：表示对应的信息维度应该几乎完全被遗忘，即在更新细胞状态时，该维度的值将被显著减小。
   • 值接近1：表示对应的信息维度应该几乎完全被保留，即在更新细胞状态时，该维度的值将基本保持不变。

   数学表达式为：

f t = σ ( W f ⋅ [ h t − 1 , x t ] + b f ) f_t = \sigma(W_f \cdot [h_{t-1}, x_t] + b_f) ft=σ(Wf⋅[ht−1,xt]+bf)

其中：

• f t f_t ft 是遗忘门在时间步 t t t 的输出向量，每个元素值都在 [ 0 , 1 ] [0, 1] [0,1] 范围内。
• σ \sigma σ 是sigmoid激活函数，它将输入值压缩到 ( 0 , 1 ) (0, 1) (0,1) 区间内，确保输出的是一个概率值。
• W f W_f Wf 是遗忘门对应的权重矩阵，用于将拼接后的输入向量映射到一个新的特征空间。
• b f b_f bf 是遗忘门的偏置项。
• [ h t − 1 , x t ] [h_{t-1}, x_t] [ht−1,xt] 表示将前一时刻隐藏状态和当前输入按维度拼接成一个单一向量。

细胞状态更新

遗忘门输出 f t f_t ft 与上一时刻细胞状态 C t − 1 C_{t-1} Ct−1 进行逐元素（element-wise）乘法，得到更新后的细胞状态 C t C_t Ct。乘积操作相当于对每个维度上的信息进行"筛选"，遗忘门输出值为0的部分会被有效遗忘（置零），为1的部分则完全保留。

C t = f t ⊙ C t − 1 C_t = f_t \odot C_{t-1} Ct=ft⊙Ct−1

这里的 ⊙ \odot ⊙ 表示逐元素乘法（Hadamard product）。通过这种方式，遗忘门能够有选择地决定哪些历史信息应当保留在细胞状态中，以便在后续时间步中继续影响模型的计算，而哪些信息应当被舍弃，从而避免无关或过时信息对当前决策产生干扰。

综上所述，LSTM的遗忘门通过计算sigmoid激活函数的输出来决定细胞状态中各维度信息的遗忘程度，并通过逐元素乘法更新细胞状态，实现了对长期依赖关系的灵活控制。这一机制使得LSTM能够在处理长序列数据时有效地避免梯度消失问题，同时保持对关键时间点信息的长期记忆能力。

输入门

LSTM（Long Short-Term Memory）网络的输入门（Input Gate）负责决定在给定时间步 t t t 时，当前输入 x t x_t xt 中哪些新信息应当被添加到细胞状态 C t C_t Ct 中。输入门的工作流程如下：

输入门计算

1. 输入合并 ：

   输入门同样接收前一时刻隐藏状态 h t − 1 h_{t-1} ht−1 和当前时间步输入 x t x_t xt，这两个向量通常被拼接（concatenate）或通过某种形式的线性变换（如全连接层）组合在一起，形成一个综合的输入向量，表示当前时间步的全部可用信息。

2. 门控值计算 ：

   综合输入向量经过两个不同的全连接层，分别生成两部分输出：

   • 输入门sigmoid部分 ：这部分输出经过sigmoid激活函数，生成一个介于0和1之间的向量 i t i_t it，表示新信息被添加到细胞状态中的比例。值接近0表示不添加，接近1表示完全添加。

   i t = σ ( W i ⋅ [ h t − 1 , x t ] + b i ) i_t = \sigma(W_i \cdot [h_{t-1}, x_t] + b_i) it=σ(Wi⋅[ht−1,xt]+bi)

   其中 W i W_i Wi 是输入门sigmoid部分的权重矩阵， b i b_i bi 是对应的偏置项， σ \sigma σ 是sigmoid激活函数。

   • 候选细胞状态（Candidate Cell State） ：这部分输出经过tanh激活函数，生成一个新的向量 C ~ t \tilde{C}_t C~t，表示可能被添加到细胞状态中的新信息。tanh函数将输出限制在 [ − 1 , 1 ] [-1, 1] [−1,1] 范围内。

   C ~ t = tanh ⁡ ( W C ⋅ [ h t − 1 , x t ] + b C ) \tilde{C}t = \tanh(W_C \cdot [h{t-1}, x_t] + b_C) C~t=tanh(WC⋅[ht−1,xt]+bC)

   其中 W C W_C WC 是候选细胞状态的权重矩阵， b C b_C bC 是对应的偏置项。

细胞状态更新

基于输入门的sigmoid输出 i t i_t it 和候选细胞状态 C ~ t \tilde{C}t C~t，通过逐元素乘法将两者结合起来，生成新信息的实际贡献值，然后将其添加到由遗忘门处理过的细胞状态 C t − 1 C{t-1} Ct−1 中，得到更新后的细胞状态 C t C_t Ct：

C t = f t ⊙ C t − 1 + i t ⊙ C ~ t C_t = f_t \odot C_{t-1} + i_t \odot \tilde{C}_t Ct=ft⊙Ct−1+it⊙C~t

这里的 ⊙ \odot ⊙ 表示逐元素乘法（Hadamard product）。通过输入门sigmoid输出 i t i_t it 与候选细胞状态 C ~ t \tilde{C}_t C~t 的乘积，模型仅允许那些被判定为重要的新信息（即 i t i_t it 中接近1的维度）进入细胞状态，并且这些新信息是以 C ~ t \tilde{C}_t C~t 中相应维度的值来更新细胞状态的。

综上所述，LSTM的输入门通过计算sigmoid激活函数的输出来决定当前输入信息中各维度应当被添加到细胞状态中的程度，并通过tanh激活函数生成候选细胞状态。然后，这两部分通过逐元素乘法结合，更新细胞状态。这样，输入门既控制了新信息的流入量，又确保了新添加的信息经过了适当的非线性变换，有助于模型捕获复杂的时间序列特征。

输出门

LSTM（Long Short-Term Memory）网络的输出门（Output Gate）负责控制在给定时间步 t t t 时，细胞状态 C t C_t Ct 中哪些信息应当被输出到当前时刻的隐藏状态 h t h_t ht 并进一步影响模型的最终输出。输出门的工作流程如下：

输出门计算

1. 输入合并 ：

   类似于输入门，输出门也接收前一时刻隐藏状态 h t − 1 h_{t-1} ht−1 和当前时间步输入 x t x_t xt，这两个向量通常通过拼接或线性变换组合成一个综合输入向量，表示当前时间步的全部可用信息。

2. 门控值计算 ：

   综合输入向量经过一个全连接层，然后使用sigmoid激活函数生成一个介于0和1之间的向量 o t o_t ot，它表示细胞状态 C t C_t Ct 中哪些信息应当被输出到隐藏状态的比例。值接近0表示不输出，接近1表示完全输出。

o t = σ ( W o ⋅ [ h t − 1 , x t ] + b o ) o_t = \sigma(W_o \cdot [h_{t-1}, x_t] + b_o) ot=σ(Wo⋅[ht−1,xt]+bo)

其中 W o W_o Wo 是输出门的权重矩阵， b o b_o bo 是对应的偏置项， σ \sigma σ 是sigmoid激活函数。

隐藏状态计算

输出门 o t o_t ot 与细胞状态 C t C_t Ct 经过tanh激活函数后的结果进行逐元素乘法，生成当前时刻的隐藏状态 h t h_t ht。tanh函数将细胞状态压缩至 [ − 1 , 1 ] [-1, 1] [−1,1] 范围内，而输出门则控制了哪些信息应当被保留并传递到后续层或作为模型的最终输出。

h t = o t ⊙ tanh ⁡ ( C t ) h_t = o_t \odot \tanh(C_t) ht=ot⊙tanh(Ct)

这里 ⊙ \odot ⊙ 仍然表示逐元素乘法。输出门 o t o_t ot 中接近1的维度对应于细胞状态 C t C_t Ct 中被选择输出的信息，接近0的维度则对应于被抑制输出的信息。通过这样的乘法操作，模型只允许那些被输出门判定为重要且应当传递给后续时间步或作为模型输出的信息，从细胞状态中流向隐藏状态。

综上所述，LSTM的输出门通过计算sigmoid激活函数的输出来决定细胞状态中各维度应当被输出到隐藏状态及后续模型计算的程度。它与经过tanh激活函数的细胞状态进行逐元素乘法，生成当前时刻的隐藏状态。这样，输出门不仅控制了细胞状态信息的流出量，还确保了输出的信息经过了适当的非线性变换，使得模型能够灵活地选择性地表达记忆细胞中的相关信息，适应不同任务的需求，如分类、回归或者生成等。

模型构建

输入

时序预测场景中，每次前向传播的input一般为[batch, input_len, node, input_channel]。

input张量的具体含义如下：

• Batch: 这个维度表示一批样本的数量。在训练或推断过程中，模型通常会一次性处理多个样本以利用并行计算的优势。

• Input Length: 这个维度通常表示时间步长、序列长度或类似的概念，即每个样本内部包含了一段连续的时间序列、事件序列或其它具有顺序关系的数据。对于每个样本，模型将在这些连续的步骤中依次处理其包含的信息。

• Node: 这个维度表示每个样本内部的节点（或称为元素、位置）数量。这里的"节点"概念通常出现在与图结构相关的数据或具有内在空间结构（如网格、点云）的数据中。每个时间步长（或序列位置）下，样本包含一组节点，每个节点都有自己的特征。

• Input Channel : 这个维度表示特征维度，即每个节点在每个时间步长（或序列位置）下具有 input_channel 个特征值。这些特征可以是数值型、类别型或其他类型的特征，共同描述了节点在当前时间步长（或序列位置）的状态。

综合来看，[batch, input_len, node, input_channel] 形状的输入表示：

• 批处理的一组样本（batch 个），每个样本包含一段具有 input_len 个时间步长（或序列位置）的数据。
• 在每个时间步长（或序列位置）下，样本内部由 node 个节点组成，每个节点具有 input_channel 个特征值。

输出

对应的output为output为[batch, output_len, node, output_channel]。

输出形状为 [batch, output_len, node, output_channel] 的四维张量表示模型对输入数据进行前向传播后的输出结果，其具体含义如下：

• Batch: 与输入相同，这个维度仍然表示一批样本的数量。模型对同一批次的样本同时进行处理，并生成对应的输出。

• Output Length : 这个维度对应于输入中的 input_len，表示模型生成的输出序列长度或时间步长。通常情况下，output_len 与 input_len 相同，意味着模型对每个输入时间步长都有相应的输出；但在某些模型（如自回归模型、循环解码器等）中，output_len 可能小于等于 input_len，尤其是当模型进行序列生成任务时，可能逐步生成输出序列。

• Node: 与输入相同，此维度表示每个样本内部的节点（或称为元素、位置）数量。模型在输出阶段同样关注这些节点，并为每个节点生成一组特征。

• Output Channel: 这个维度表示模型为每个节点在每个输出时间步长（或序列位置）生成的特征维度。每个输出特征通道可能对应于某种特定的预测结果、概率分布、注意力权重、状态变量等，具体取决于模型的设计和应用任务。

综上所述，形状为 [batch, output_len, node, output_channel] 的输出张量表示：

• 批处理的一组样本（batch 个），每个样本生成一段具有 output_len 个时间步长（或序列位置）的输出数据。
• 在每个输出时间步长（或序列位置）下，样本内部的每个节点都有 output_channel 个特征值，这些特征值反映了模型对节点在该时刻状态的预测、解释或生成结果。

模型适配

要基于PyTorch构建一个能够处理输入形状为 [batch, input_len, node, input_channel] 并产生输出形状为 [batch, output_len, node, output_channel] 的LSTM模型，首先需要明确LSTM层本身并不直接支持这种四维输入/输出。LSTM通常处理的是三维张量，即 [batch_size, sequence_length, input_features] 的输入以及 [batch_size, sequence_length, hidden_size] 或 [batch_size, sequence_length, num_directions * hidden_size]（双向LSTM）的输出。因此，需要适当调整模型结构或对数据进行预处理，以适应LSTM的要求。

以下是常用的处理方式：

• 展平节点和通道信息：

  如果输入数据中每个节点在每个时间步长有多个特征（即input_channel），那么可以先将这些特征展平到一维，使得每个时间步长每个节点只有一个值。类似地，对于输出，也需要将output_channel展平。这样处理后，输入和输出的形状分别变为 [batch, input_len, node * input_channel] 和 [batch, output_len, node * output_channel]。

如果本身的输入和输出的channel均为1，比如基于最近12个时间窗的流量预测未来3个时间窗的流量，则直接可将[batch, input_len, node, 1]通过squeeze函数直接变为[batch, input_len, node]，即：

input = input.squeeze(-1)

模型开发

基于上述内容，构建模型如下：

class LSTM(nn.Module):
    def __init__(self, node_num, input_len, input_channel, hidden_sizes, output_len, output_channel):
        super(LSTM, self).__init__()
        self.node_num = node_num
        self.input_len = input_len
        self.input_channel = input_channel
        self.hidden_sizes = hidden_sizes
        self.output_len = output_len
        self.layers = nn.ModuleList()
        self.output_channel = output_channel
        prev_size = node_num * input_channel
        for hidden_size in hidden_sizes:
            self.layers.append(nn.LSTM(input_size=prev_size, hidden_size=hidden_size))
            prev_size = hidden_size
        self.layers.append(nn.LSTM(input_size=prev_size, hidden_size=node_num * output_channel))
        self.mlp = nn.Linear(input_len, output_len)
    def forward(self, x):
        # [Batch, Input_len, Node, Input_channel] --> [Batch, Input_len, Node*Input_channel]
        x = x.reshape(x.shape[0], x.shape[1], -1)
        # [Batch, Input_len, Node*Input_channel] --> [Input_len, Batch, Node*Input_channel]
        x = x.permute(1, 0, 2)
        prev_output = x
        for layer in self.layers:
            output, (ht, ct) = layer(prev_output)
            prev_output = output
        # [Input_len, Batch, Node * Output_channel] --> [Batch, Node * Output_channel, Input_len]
        y = output.permute(1, 2, 0)
        # [Batch, Node * Output_channel, Input_len] --> [Batch, Node * Output_channel, Output_len]
        y = self.mlp(y)
        # [Batch, Node * Output_channel, Output_len] --> [Batch, Output_len, Node * Output_channel]
        y = y.permute(0, 2, 1)
        y = y.reshape(y.shape[0], y.shape[1], -1, self.output_channel)
        return y

上述代码定义了一个名为 LSTM 的 PyTorch 模块类，继承自 nn.Module。该类旨在处理形状为 [Batch, Input_len, Node, Input_channel] 的输入数据，并生成形状为 [Batch, Output_len, Node, Output_channel] 的输出。

初始化方法 (__init__)：

• 类接收6个参数：node_num（节点数）、input_len（输入序列长度）、input_channel（输入通道数）、hidden_sizes（隐藏层大小列表）、output_len（输出序列长度）和 output_channel（输出通道数）。

• 初始化父类 nn.Module。

• 将传入的参数存储为类属性。

• 定义一个可迭代的模块列表 self.layers，用于存放多层 LSTM 单元。

• 初始化一个 prev_size 变量，其初始值为 node_num * input_channel，表示展平后的节点特征总维度。

• 遍历 hidden_sizes 列表，为每层 LSTM 定义一个 nn.LSTM 单元，其中 input_size 设置为当前 prev_size，hidden_size 设置为当前遍历到的隐藏层大小。每次迭代后，更新 prev_size 为当前隐藏层大小，以便下一层 LSTM 使用。

• 添加最后一层 LSTM，其 input_size 为上一层 LSTM 的输出维度（即最后一个 hidden_size），hidden_size 为 node_num * output_channel。这样最后一层 LSTM 的输出可以直接映射到所需形状的输出。

• 定义一个全连接层（MLP）self.mlp，其输入维度为 input_len，输出维度为 output_len。该层用于将 LSTM 的输出时间步长调整为所需的 output_len。

前向传播方法 (forward)：

• 输入 x 形状为 [Batch, Input_len, Node, Input_channel]。

• 展平节点和通道信息 ：将输入 reshape 为 [Batch, Input_len, Node * Input_channel]，将节点和通道特征合并为一维。

• 调整输入顺序 ：将输入 permute 为 [Input_len, Batch, Node * Input_channel]，使时间步长成为第一维，便于 LSTM 处理。

• 初始化 prev_output 为经过调整顺序后的输入 x，用于存储当前层的输出，供下一层 LSTM 使用。

• 逐层处理 ：遍历 self.layers 中的所有 LSTM 单元，对 prev_output 进行前向传播，得到当前层的输出和隐含状态（output, (ht, ct)）。将当前输出赋值给 prev_output，准备传递给下一层。

• 调整输出顺序 ：将最后一层 LSTM 的输出 permute 回 [Batch, Node * Output_channel, Input_len]，恢复到原始的批次和节点特征维度顺序，时间步长为第三维。

• 调整时间步长 ：使用全连接层 self.mlp 对 LSTM 输出进行前向传播，将时间步长从 Input_len 调整为 Output_len，得到形状为 [Batch, Node * Output_channel, Output_len] 的中间结果。

• 重新排列维度 ：将中间结果 permute 为 [Batch, Output_len, Node * Output_channel]，使得输出序列长度成为第二维。

• 重塑为所需输出形状 ：最后，将结果 reshape 为 [Batch, Output_len, Node, Output_channel]，满足预期输出形状。

• 返回处理后的输出。

总结来说，上述LSTM 类实现了对形状为 [Batch, Input_len, Node, Input_channel] 的输入数据进行多层 LSTM 处理，并通过全连接层调整时间步长，最终生成形状为 [Batch, Output_len, Node, Output_channel] 的输出。在处理过程中，节点特征被展平并与通道特征合并，以适应 LSTM 的输入要求。

训练、验证、测试

训练、验证、测试的代码可直接参见基于MLP算法实现交通流量预测(Pytorch版)。

代码基本一致，只需要修改Exp_LSTM_PEMS的_build_model即可：

def _build_model(self):
	if self.args.dataset == 'PEMS03':
		self.input_dim = 358
	elif self.args.dataset == 'PEMS04':
		self.input_dim = 307
	elif self.args.dataset == 'PEMS07':
		self.input_dim = 883
	elif self.args.dataset == 'PEMS08':
		self.input_dim = 170
	model = LSTM(node_num=self.input_dim, input_len=args.window_size, input_channel=1, hidden_sizes=args.hidden_sizes,
			 output_len=args.horizon, output_channel=1)
	print(model)
	return model

上述模型，输入仅使用了流量这1个特征，然后来预测流量这1个特征，所以input_channel和output_channel均为1，若使用多变量预测单变量，或者多变量预测多变量，读者可自行调整配置，此处不再赘述。

模型效果

最后，将我们构建好的MLP网络在PEMS数据集进行了准确性测试，算法测试的相关配置如下:

torch.manual_seed(4321)  # reproducible
parser = argparse.ArgumentParser(description='LSTM on pems datasets')
### -------  dataset settings --------------
parser.add_argument('--dataset', type=str, default='PEMS08',
					choices=['PEMS03', 'PEMS04', 'PEMS07', 'PEMS08'])  # sometimes use: PeMS08
parser.add_argument('--norm_method', type=str, default='z_score')
parser.add_argument('--normtype', type=int, default=0)
### -------  input/output length settings --------------
parser.add_argument('--window_size', type=int, default=12)
parser.add_argument('--horizon', type=int, default=12)
parser.add_argument('--train_length', type=float, default=6)
parser.add_argument('--valid_length', type=float, default=2)
parser.add_argument('--test_length', type=float, default=2)
### -------  training settings --------------
parser.add_argument('--use_gpu', type=bool, default=False)
parser.add_argument('--train', type=bool, default=True)
parser.add_argument('--resume', type=bool, default=False)
parser.add_argument('--evaluate', type=bool, default=False)
parser.add_argument('--finetune', type=bool, default=False)
parser.add_argument('--validate_freq', type=int, default=1)
parser.add_argument('--epoch', type=int, default=80)
parser.add_argument('--lr', type=float, default=0.001)
parser.add_argument('--batch_size', type=int, default=8)
parser.add_argument('--optimizer', type=str, default='N')  #
parser.add_argument('--early_stop', type=bool, default=True)
parser.add_argument('--early_stop_step', type=int, default=5)
parser.add_argument('--exponential_decay_step', type=int, default=5)
parser.add_argument('--decay_rate', type=float, default=0.5)
parser.add_argument('--lradj', type=int, default=1, help='adjust learning rate')
parser.add_argument('--weight_decay', type=float, default=1e-5)
parser.add_argument('--model_name', type=str, default='LSTM')
### -------  model settings --------------
parser.add_argument('--hidden_sizes', type=list, default=[64, 36, 64])
args = parser.parse_args()
Exp=Exp_LSTM_PEMS
exp=Exp(args)

可以看到，我们定义了4层LSTM，以PEMS03为例，因为PEMS03的节点数为358，所以4层LSTM的维度变化如下:

第1个LSTM: [batch, input_len, 358] --> [batch, input_len, 64]
第2个LSTM: [batch, input_len, 64] --> [batch, input_len, 36]
第3个LSTM: [batch, input_len, 36] --> [batch, input_len, 64]
第4个LSTM: [batch, input_len, 64] --> [batch, input_len, 358]

结果如下:

Dataset	MAE	MAPE	RMSE
PEMS03	19.1957	0.181804	32.8948
PEMS04	23.2434	0.155689	37.2810
PEMS07	29.7381	0.129077	47.9241
PEMS08	20.4671	0.127739	33.2526