| 33 | """ |
| 34 | |
| 35 | def __init__(self, input_size, hidden_size, num_layers=1, |
| 36 | nonlinearity="tanh", bias=True, batch_first=False, dropout=0, |
| 37 | bidirectional=False, rnn_type=RNNType.GRU): |
| 38 | super(RNN, self).__init__() |
| 39 | self.rnn_type = rnn_type |
| 40 | self.num_layers = num_layers |
| 41 | self.batch_first = batch_first |
| 42 | self.bidirectional = bidirectional |
| 43 | if rnn_type == RNNType.LSTM: |
| 44 | self.rnn = torch.nn.LSTM( |
| 45 | input_size, hidden_size, num_layers=num_layers, bias=bias, |
| 46 | batch_first=batch_first, dropout=dropout, |
| 47 | bidirectional=bidirectional) |
| 48 | elif rnn_type == RNNType.GRU: |
| 49 | self.rnn = torch.nn.GRU( |
| 50 | input_size, hidden_size, num_layers=num_layers, bias=bias, |
| 51 | batch_first=batch_first, dropout=dropout, |
| 52 | bidirectional=bidirectional) |
| 53 | elif rnn_type == RNNType.RNN: |
| 54 | self.rnn = torch.nn.RNN( |
| 55 | input_size, hidden_size, vnonlinearity=nonlinearity, bias=bias, |
| 56 | batch_first=batch_first, dropout=dropout, |
| 57 | bidirectional=bidirectional) |
| 58 | else: |
| 59 | raise TypeError( |
| 60 | "Unsupported rnn init type: %s. Supported rnn type is: %s" % ( |
| 61 | rnn_type, RNNType.str())) |
| 62 | |
| 63 | def forward(self, inputs, seq_lengths=None, init_state=None, |
| 64 | ori_state=False): |