Feature Request: Support LayerNormalization in TFLite

[christian-steinmeyer]

As integral part of many transformer models, it'd be great if armnn supported LayerNormalization.

[TeresaARM]

Hello @christian-steinmeyer Arm NN has a Normalization layer /armnn/src/armnn/layers/NormalizationLayer.cpp

Which error are you seeing? Maybe we are missing some parameter?

[christian-steinmeyer]

Hi! To add a bit of context, I'm working on pruning a tensorflow model and using xnnpack to accelerate inference with it on Android. To that end, I have a model (see below), convert it to tflite using tf.lite.Optimize.EXPERIMENTAL_SPARSITY and run the tensorflow benchmark tool on it.

Here's the toy model I'm using:

model = keras.Sequential(
    [
        keras.layers.InputLayer(input_shape=(28, 28, 1)),
        keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation='relu'),
        keras.layers.MaxPooling2D(pool_size=(2, 2)),
        keras.layers.Conv2D(1024, (7,7)),
        keras.layers.Flatten(),
        keras.layers.LayerNormalization(),  # works without, but not with this normalization layer
        keras.layers.Dense(10, activation="softmax"),
    ]
)
# ...
model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(model, **pruning_params)
model_for_pruning.fit(
    # ...
)
pruned_model = tfmot.sparsity.keras.strip_pruning(model_for_pruning)

converter = tf.lite.TFLiteConverter.from_keras_model(pruned_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT, tf.lite.Optimize.EXPERIMENTAL_SPARSITY]
converter.target_spec.supported_types = [tf.float16]
converter.allow_custom_ops = True
tflite_model = converter.convert()

When I include the LayerNormalization layer, running the benchmark tool, I get the following log output: tflite : Failed to apply XNNPACK delegate. vs Replacing 12 out of 12 node(s) with delegate (TfLiteXNNPackDelegate) node, yielding 1 partitions for the whole graph without this layer.

I use this command to run the benchmark: adb shell taskset f0 am start -S -n org.tensorflow.lite.benchmark/.BenchmarkModelActivity --es args '"--verbose=true --num_threads=1 --warmup_runs=100 --num_runs=100 --use_xnnpack=true --graph=/data/local/tmp/model.tflite"'

[MikeJKelly]

Hi @christian-steinmeyer

the XNNPACK delegate is one of the TensorFlow delegates, (see https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/delegates/xnnpack ) not the ArmNN delegate.

From looking at the XNNPACK page it doesn't seem like they support LayerNormalization yet but they may be able to help you there.

Best regards, Mike

[christian-steinmeyer]

Thanks Mike! I thought somehow armnn was responsible. Thanks for the pointer. I think this belongs to the tensorflow repo then. Best, Chris