LSTM/GRU详细代码解析+完整代码实现

LSTM和GRU目前被广泛的应用在各种预测场景中，并与卷积神经网络CNN或者图神经网络GCN这里等相结合，对数据的结构特征和时序特征进行提取，从而预测下一时刻的数据。在这里整理一下详细的LSTM/GRU的代码，并基于heatmap热力图实现对结果的展示。

一、GRU

GRU的公式如下图所示：

 其代码部分：

class GRU(torch.nn.Module):
    def __init__(self, hidden_size, output_size, num_layers):
        super().__init__()
        self.input_size = 1
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.output_size = output_size
        self.num_directions = 1
        self.gru = torch.nn.GRU(self.input_size, self.hidden_size, self.num_layers, batch_first=True)
        self.linear = torch.nn.Linear(self.hidden_size, self.output_size)

    def forward(self, input_seq):
        # input(batch_size, seq_len, input_size)
        batch_size, seq_len = input_seq.shape[0], input_seq.shape[1]
        h_0 = torch.randn(self.num_directions * self.num_layers, batch_size, self.hidden_size).to(device)
        # output(batch_size, seq_len, num_directions * hidden_size)
        output, _ = self.gru(input_seq, (h_0))
        pred = self.linear(output)
        pred = pred[:, -1, :]
        return pred

这里主要对里面主要的五个参数进行介绍： 

 input_size：输入节点特征的维度。这里需要注意的是，如果你输入的是节点的交通流量数据，一般只使用一个值表示，那么你的input_size为1；若是想基于该节点在t时刻的多个特征，如：流量、速度、车辆数这三个指标对交通流量进行预测，这是input_size=3。

hidden_size