Pytorch GRU

import torch
import torch.nn as nn


class GRU(nn.Module):
    '''GRU + 全连接'''
    def __init__(self, num_layers=1, input_size=8,
                 hidden_size=64, time_step=20, output_size=1):
        super(GRU, 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.GRU = nn.GRU(num_layers=num_layers, input_size=input_size,
                          hidden_size=hidden_size, batch_first=True)  # batch优先

        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)
        '''
        gru_out, _ = self.GRU(x)
        out = self.linear(gru_out)
        return out[:, -1, :] # 返回最后一个时间步 (batch_size, output_size)



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