<p>Based on this <a href="https://arxiv.org/abs/1607.06450" rel="nofollow noreferrer">paper</a>: <em>"Layer Normalization" - Jimmy Lei Ba, Jamie Ryan Kiros, Geoffrey E. Hinton</em></p>

<p>Tensorflow now comes with the <a href="https://github.com/tensorflow/tensorflow/blob/r1.12/tensorflow/contrib/rnn/python/ops/rnn_cell.py" rel="nofollow noreferrer"><code>tf.contrib.rnn.LayerNormBasicLSTMCell</code></a> a LSTM unit with layer normalization and recurrent dropout.</p>

<p>Find the documentation <a href="https://www.tensorflow.org/api_docs/python/tf/contrib/rnn/LayerNormBasicLSTMCell" rel="nofollow noreferrer">here</a>.</p>
How to implement Tensorflow batch normalization in LSTM
<p>My current LSTM network looks like this.</p>

<pre class="lang-py prettyprint-override"><code>rnn_cell = tf.contrib.rnn.BasicRNNCell(num_units=CELL_SIZE)
init_s = rnn_cell.zero_state(batch_size=1, dtype=tf.float32)  # very first hidden state
outputs, final_s = tf.nn.dynamic_rnn(
    rnn_cell,              # cell you have chosen
    tf_x,                  # input
    initial_state=init_s,  # the initial hidden state
    time_major=False,      # False: (batch, time step, input); True: (time step, batch, input)
)

# reshape 3D output to 2D for fully connected layer
outs2D = tf.reshape(outputs, [-1, CELL_SIZE])
net_outs2D = tf.layers.dense(outs2D, INPUT_SIZE)

# reshape back to 3D
outs = tf.reshape(net_outs2D, [-1, TIME_STEP, INPUT_SIZE])
</code></pre>

<p>Usually, I apply <a href="https://www.tensorflow.org/api_docs/python/tf/layers/batch_normalization" rel="noreferrer"><code>tf.layers.batch_normalization</code></a> as batch normalization. But I am not sure if this works in a LSTM network.</p>

<pre class="lang-py prettyprint-override"><code>b1 = tf.layers.batch_normalization(outputs, momentum=0.4, training=True)
d1 = tf.layers.dropout(b1, rate=0.4, training=True)

# reshape 3D output to 2D for fully connected layer
outs2D = tf.reshape(d1, [-1, CELL_SIZE])                       
net_outs2D = tf.layers.dense(outs2D, INPUT_SIZE)

# reshape back to 3D
outs = tf.reshape(net_outs2D, [-1, TIME_STEP, INPUT_SIZE])
</code></pre>