Apply Dense to every output in RNN Layer

Viewed 1157

I am trying to convert the following Keras code into pure Tensorflow but I have trouble adding a dense layer to every timestep of the bidirectional RNN output:

Here is the Keras code in question:

self.model = Sequential()
self.model.add(Bidirectional(LSTM(nr_out, return_sequences=True,
                                     dropout_W=dropout, dropout_U=dropout),
                                     input_shape=(max_length, nr_out)))
self.model.add(TimeDistributed(Dense(nr_out, activation='relu', init='he_normal')))
self.model.add(TimeDistributed(Dropout(0.2)))

Here is the initial tensorflow code:

lstm_cell_fwd = rnn.BasicLSTMCell(num_hidden, forget_bias=1.0)
lstm_cell_bwd = rnn.BasicLSTMCell(num_hidden, forget_bias=1.0)
outputs, output_state_fw, output_state_bw  = rnn.static_bidirectional_rnn(lstm_cell_fwd, lstm_cell_bwd, inputs=sequence, dtype=tf.float64)

Generally if I only wanted to predict on the last state I would do something like:

logits = tf.matmul(outputs[-1], weights['out']) + biases['out']

What is the best way to express the TimeDistributed layer in Tensorflow?

1 Answers

Try updating your cell definition to something like:

lstm_cell_fwd = rnn.BasicLSTMCell(num_hidden, forget_bias=1.0)
lstm_cell_fwd = rnn.DropoutWrapper(lstm_cell_fwd, input_keep_prob=dropout, output_keep_prob=dropout)  # if you want to keep dropout, which seems to be in your Keras model
lstm_cell_fwd = rnn.OutputProjectionWrapper(lstm_cell_fwd, nr_out)  # FC output layer
# Similarly for lstm_cell_bwd
outputs, output_state_fw, output_state_bw  = rnn.static_bidirectional_rnn(lstm_cell_fwd, lstm_cell_bwd, ...)

It looks like your Keras definition uses dropout, so I added a dropout layer here. I believe the dropout_W from Keras is equivalent to input_keep_prob in TF, and dropout_U from Keras is output_keep_prob in TF. For your dropout layer, you would need to define a placeholder:

dropout = tf.placeholder(tf.float32, [], name='dropout')

and feed that when you run your network with some dropout probability for training, and usually dropout=1.0 for validation, testing, and using the network.

Related