Pytorch LSTM

import torch
import torch.nn as nn



class LSTM(nn.Module):
    '''LSTM + 全连接'''
    def __init__(self, num_layers=1, input_size=8,
                 hidden_size=64, time_step=20, output_size=1):
        super(LSTM, self).__init__()

        self.num_layers, self.input_size, self.hidden_size, self.time_step, self.output_size = \
            num_layers, input_size, hidden_size, time_step, output_size

        self.LSTM = nn.LSTM(num_layers=num_layers, input_size=input_size,
                            hidden_size=hidden_size, batch_first=True)

        self.linear = nn.Linear(hidden_size * num_layers, output_size)

    def forward(self, x):
        '''
        :param x: (batch_size, time_step, input_size) batch_first=True
        :return: (batch_size, output_size)
        '''
        lstm_out, (h_n, h_c) = self.LSTM(x, None)  # None表示hidden_state会用全0的state
        out = self.linear(lstm_out)
        return out[:, -1, :] # 返回最后一个时间步(batch_size, output_size)



if __name__ == '__main__':
    data = torch.randn(1600).view(10, 20, 8)
    model = LSTM(output_size=2)
    out = model(data)
    print(data.size()) # torch.Size([10, 20, 8])
    print(out.size()) # torch.Size([10, 2])