Class RNN

model/rnn.py:30–108

One layer rnn.

class RNN(torch.nn.Module):
    """
    One layer rnn.
    """

    def __init__(self, input_size, hidden_size, num_layers=1,
                 nonlinearity="tanh", bias=True, batch_first=False, dropout=0,
                 bidirectional=False, rnn_type=RNNType.GRU):
        super(RNN, self).__init__()
        self.rnn_type = rnn_type
        self.num_layers = num_layers
        self.batch_first = batch_first
        self.bidirectional = bidirectional
        if rnn_type == RNNType.LSTM:
            self.rnn = torch.nn.LSTM(
                input_size, hidden_size, num_layers=num_layers, bias=bias,
                batch_first=batch_first, dropout=dropout,
                bidirectional=bidirectional)
        elif rnn_type == RNNType.GRU:
            self.rnn = torch.nn.GRU(
                input_size, hidden_size, num_layers=num_layers, bias=bias,
                batch_first=batch_first, dropout=dropout,
                bidirectional=bidirectional)
        elif rnn_type == RNNType.RNN:
            self.rnn = torch.nn.RNN(
                input_size, hidden_size, vnonlinearity=nonlinearity, bias=bias,
                batch_first=batch_first, dropout=dropout,
                bidirectional=bidirectional)
        else:
            raise TypeError(
                "Unsupported rnn init type: %s. Supported rnn type is: %s" % (
                    rnn_type, RNNType.str()))

    def forward(self, inputs, seq_lengths=None, init_state=None,
                ori_state=False):
        """
        Args:
            inputs:
            seq_lengths:
            init_state:
            ori_state: If true, will return ori state generate by rnn. Else will
                       will return formatted state
        :return:
        """
        if seq_lengths is not None:
            seq_lengths = seq_lengths.int()
            sorted_seq_lengths, indices = torch.sort(seq_lengths,
                                                     descending=True)
            if self.batch_first:
                sorted_inputs = inputs[indices]
            else:
                sorted_inputs = inputs[:, indices]
            packed_inputs = torch.nn.utils.rnn.pack_padded_sequence(
                sorted_inputs, sorted_seq_lengths.cpu(), batch_first=self.batch_first)
            outputs, state = self.rnn(packed_inputs, init_state)
        else:
            outputs, state = self.rnn(inputs, init_state)