upgrade to tf2

Hi, 
I tried to adjust the code in order to work with TF2. However I'm new in TF and I got stuck and I could not find the alternative for `tf.contrib.layers.layer_norm` in TF2, some people said that `tf.keras.layers.LayerNormalization` will be equivalent but input parameters are different. 

Thanks for your hardworking, 
I'm looking forward to seeing the code is compatible on TF2 