Translating LSTM implementation from tensorflow to pytorch

Viewed 94

I'm attemping to translate an old TensorFlow code into a PyTorch equivalent, and I'm struggling with understanding TensorFlow's implementation of LSTMs.

Original code:

        rnn_cell_basic = tf.nn.rnn_cell.BasicLSTMCell(self.rnn_size, state_is_tuple=False)
        cell = tf.nn.rnn_cell.MultiRNNCell([rnn_cell_basic] * self.num_rnn_layers, state_is_tuple=False)

How does that translate into PyTorch?

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

My guess is num_layers = num_rnn_layers, however as for input_size and hidden_size, I'm not sure which one's which.

0 Answers
Related