定义LSTM 模型

定义LSTM神经网络模型

class LSTM(nn.Module):
    """
        Parameters：
        - input_size:  输入特征数
        - hidden_size:   单个隐藏层的节点数
        - output_size: 输出特征数
        - num_layers:  隐藏层数
    """
    def __init__(self, input_size=1, hidden_size=1, output_size=1, num_layers=1):
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.output_size = output_size
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers)
        self.forwardCalculation = nn.Linear(hidden_size, output_size)
        self.hidden_cell = (torch.zeros(self.num_layers,self.output_size,self.hidden_size),
                            torch.zeros(self.num_layers,self.output_size,self.hidden_size))

    def forward(self, input_seq):
        L = len(input_seq)
        lstm_out, self.hidden_cell = self.lstm(input_seq.view(L, 1, -1),self.hidden_cell)
        predictions = self.forwardCalculation(lstm_out.view(len(input_seq), -1))
        return predictions[-1]