Can't convert tensorflow.keras.layers.LayerNormalization

[diegodgs]

Hello,

I am using LayerNormalization on my models to replace BacthNormalization in some scenarios. When I try to convert to onnx, a keras model that uses LayerNormalization. I receive the following error.

ValueError: Unable to find out a correct type for tensor type = 20 of layer_normalization/add/ReadVariableOp/resource:0

I am using tensorflow 2.2 and tf.keras.

[jiafatom]

LayerNormalization is not supported yet, can you provide a simple example for us to debug?

[diegodgs]

Hello, Thank you for your time, I create the following simple example, Its possible to reproduce and debug. How I could learn about keras2onnx and try to help?

import tensorflow as tf
import keras2onnx


model = tf.keras.Sequential(
    [
            tf.keras.layers.Input((224, 224, 3)),
            tf.keras.layers.Conv2D(8, 3),
            tf.keras.layers.LayerNormalization(),
            tf.keras.layers.LeakyReLU(),
        ]
)
model.summary()

onnx_model = keras2onnx.convert_keras(model, model.name)
keras2onnx.save_model(onnx_model, 'model.onnx')

[jiafatom]

Thanks for your interest trying to help. Here you can see a map keras_layer_to_operator from keras layer to conversion function. So you just need create a conversion function for LayerNormalization layer. Check other conversion functions, then you can get an idea. The conversion function is used to construct this keras layer by a few onnx operators. This is the key for your issue.

[diegodgs]

Hello, starting some study about this problem and finding a way to help.

First, I start to visit the onnx operators. I am inclined to believe that InstanceNormalization operator do the same thing as LayerNormalization. It's correct? If the answers is yes. Could be simple to provide LayerNormalization support. What is your opinion @jiafatom ?

[jiafatom]

See here an illustration of different normalization. I feel that layer normalization is similar to batch normalization, just need add Transpose to switch axis and then Transpose back.

[bwery]

I encounter the same problem. So it would be nice if this operator can be supported in future !

[bwery]

@jiafatom: Having a look to your idea to use Transpose to modify the behaviour of a batch normalization layer into the behaviour of a layer normalization one, I feel there is a problem with inference.

I understand that batch normalization layers are supposed to be frozen during inference, working with long term statistically derived parameters, while layer normalization is not. This is because the behaviour of the layer normalization is determined by the content of the current set of data only, and therefore will have no "history effect".

[marcusturewicz]

This has shown up for me as an issue when trying to convert the EfficientNet models to ONNX. I'm happy to contribute to supporting LayerNormalization.

[marcusturewicz]

I've found a workaround. You can export the Keras model to TensorFlow's SavedModel format, then use tf2onnx to convert from SavedModel to ONNX. For example:

Install things:

pip install tensorflow tf2onnx

Run the following Python script to save the SaveModel format to disk:

import tensorflow as tf

saved_model_dir = "efficientnetb3"
tf_keras_model = tf.keras.applications.EfficientNetB3()
tf.saved_model.save(tf_keras_model, saved_model_dir)

Then run the following from the command line to convert to onnx:

python -m tf2onnx.convert --saved-model efficientnetb3 --opset 10 --output efficientnetb3.onnx --fold_const

And after a while:

<<Lots of logs omitted for brevity>>
2020-11-18 16:01:10,295 - INFO - Using tensorflow=2.3.1, onnx=1.7.0, tf2onnx=1.6.2/8d5253
2020-11-18 16:01:10,296 - INFO - Using opset <onnx, 10>
2020-11-18 16:01:21,756 - INFO - Optimizing ONNX model
2020-11-18 16:02:04,619 - INFO - After optimization: BatchNormalization -71 (78->7), Const -483 (790->307), Identity -25 (25->0), Squeeze -26 (26->0), Transpose -414 (442->28), Unsqueeze -104 (104->0)
2020-11-18 16:02:04,784 - INFO -
2020-11-18 16:02:04,784 - INFO - Successfully converted TensorFlow model C:\\onnxdemo\\tensorflow\\efficientnetb3 to ONNX
2020-11-18 16:02:05,009 - INFO - ONNX model is saved at C:\\onnxdemo\\tensorflow\\efficientnetb3.onnx