Data-Augmentation layers can't be exported to ONNX

[lambda-science]

Hi,

model.export(os.path.join(MODEL_PATH, "model.onnx"), format="onnx", input_signature=[keras.InputSpec(shape=(None, 256, 256, 3), dtype="float32")])

Export a model to ONNX format but with a lot of error and warnings. Such as:

ERROR:tf2onnx.tfonnx:Tensorflow op [sequential_2_1/sequential_1_1/random_brightness_1_1/stateless_random_uniform/StatelessRandomUniformV2: StatelessRandomUniformV2] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [sequential_2_1/sequential_1_1/random_contrast_1_1/stateless_random_uniform/StatelessRandomUniformV2: StatelessRandomUniformV2] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [sequential_2_1/sequential_1_1/random_flip_1_1/stateless_random_uniform/StatelessRandomUniformV2: StatelessRandomUniformV2] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [sequential_2_1/sequential_1_1/random_rotation_1_1/stateless_random_uniform/StatelessRandomUniformV2: StatelessRandomUniformV2] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [sequential_2_1/sequential_1_1/random_rotation_1_1/ImageProjectiveTransformV3: ImageProjectiveTransformV3] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [sequential_2_1/sequential_1_1/random_zoom_1_1/stateless_random_uniform_1/StatelessRandomUniformV2: StatelessRandomUniformV2] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [sequential_2_1/sequential_1_1/random_zoom_1_1/stateless_random_uniform/StatelessRandomUniformV2: StatelessRandomUniformV2] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [sequential_2_1/sequential_1_1/random_zoom_1_1/ImageProjectiveTransformV3: ImageProjectiveTransformV3] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [sequential_2_1/sequential_1_1/random_translation_1_1/stateless_random_uniform_1/StatelessRandomUniformV2: StatelessRandomUniformV2] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [sequential_2_1/sequential_1_1/random_translation_1_1/stateless_random_uniform/StatelessRandomUniformV2: StatelessRandomUniformV2] is not supported
ERROR:tf2onnx.tfonnx:Tensorflow op [sequential_2_1/sequential_1_1/random_translation_1_1/ImageProjectiveTransformV3: ImageProjectiveTransformV3] is not supported
ERROR:tf2onnx.tfonnx:Unsupported ops: Counter({'StatelessRandomUniformV2': 8, 'ImageProjectiveTransformV3': 3})

And then when I try to import it in ONNX runtime I get the following erros:

InvalidGraph: [ONNXRuntimeError] : 10 : INVALID_GRAPH : Load model from model.onnx failed:This is an invalid model. In Node, 
("sequential_2_1/sequential_1_1/random_brightness_1_1/stateless_random_uniform/StatelessRandomUniformV2", StatelessRandomUniformV2, "", -1) : 
("sequential_2_1/sequential_1_1/random_brightness_1_1/stateless_random_uniform/shape_Concat__88:0": tensor(int32),"ConstantFolding/sequential_2_1/sequential_1_1/random_brightness_1_1/stateless_random_uniform/StatelessRandomGetKeyCounter-folded-0:0": 
tensor(uint64),"ConstantFolding/sequential_2_1/sequential_1_1/random_brightness_1_1/stateless_random_uniform/StatelessRandomGetKeyCounter-folded-1:0": 
tensor(uint64),"sequential_2_1/sequential_1_1/random_translation_1_1/stateless_random_uniform_1/StatelessRandomUniformV2/alg:0": tensor(int32),) -> 
("sequential_2_1/sequential_1_1/random_brightness_1_1/stateless_random_uniform/StatelessRandomUniformV2:0",) , Error No Op registered for StatelessRandomUniformV2 with domain_version of 15

My model code is something like

data_augmentation = keras.Sequential([
    keras.layers.RandomBrightness(factor=0.2, input_shape=(None, None, 3)),
    keras.layers.RandomContrast(factor=0.2),
    keras.layers.RandomFlip("horizontal_and_vertical"),
    keras.layers.RandomRotation(0.3, fill_mode="constant"),
    keras.layers.RandomZoom(.2, .2, fill_mode="constant"),
    keras.layers.RandomTranslation(0.2, .2, fill_mode="constant"),
    keras.layers.Resizing(256, 256, interpolation="bilinear", crop_to_aspect_ratio=True),
    keras.layers.Rescaling(scale=1. / 127.5, offset=-1),  # For [-1, 1] scaling
])

# My ResNet50V2
model = keras.models.Sequential()
model.add(data_augmentation)
model.add(
    ResNet50V2(
        include_top=False,
        input_shape=(256, 256, 3),
        pooling="avg",
    )
)
model.add(keras.layers.Flatten())
model.add(keras.layers.Dense(len(config.SUB_FOLDERS), activation='softmax'))

model.summary()

model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=config.MAX_LR),
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=False),
    metrics=[config.METRIC]
)

My environement informations:

• MacOS

[[package]]
name = "tensorflow"
version = "2.19.0"
[[package]]
name = "tensorflow-metal"
version = "1.2.0"
[[package]]
name = "keras"
version = "3.10.0"
[[package]]
name = "numpy"
version = "2.1.3"
[[package]]
name = "tf2onnx"
version = "1.16.1"
[[package]]
name = "onnx"
version = "1.17.0"

[lambda-science]

Update: I removed the data augmentation layers so it works but I get widly different results between the model in ONNX version or Keras version.

import keras
import numpy as np
import datasets
import onnxruntime as ort
import tensorflow as tf

model = keras.models.load_model("model.keras")

# Load image from HuggingFace
ds = datasets.load_dataset("corentinm7/MyoQuant-SDH-Data")

# Load ONNX model
session = ort.InferenceSession("model.onnx")
input_name = session.get_inputs()[0].name

# TensorFlow SavedModel inference
saved_model = tf.saved_model.load("tf_saved_format")
infer = saved_model.signatures["serving_default"]  # Default inference signature
image = ds["test"][3]["image"]
image_batch = np.expand_dims(np.array(image, dtype="float32"), axis=0)
# Shape: (1, H, W, 3)# PIL image
# Inference
outputs = session.run(None, {input_name: image_batch})
# Get softmax output
probs = outputs[0][0]  # shape: (2,)

# Predicted class index
predicted_class = np.argmax(probs)

# Confidence score
confidence = probs[predicted_class]

print(f"ONNX Predicted class: {predicted_class}, confidence: {confidence:.2f}")


test_proba = model.predict(image_batch)
test_classes = test_proba.argmax(axis=-1)

print(f"KERAS Predicted class: {test_classes}, confidence: {max(test_proba)}")

# Convert numpy to tf.Tensor
image_tensor = tf.convert_to_tensor(image_batch, dtype=tf.float32)
tf_output = infer(image_tensor)
# Get output tensor name (usually only one output)
output_tensor = list(tf_output.values())[0].numpy()[0]
tf_pred = np.argmax(output_tensor)
tf_conf = output_tensor[tf_pred]

print(f"TF SavedModel Predicted class: {tf_pred}, confidence: {tf_conf:.2f}")

ONNX Predicted class: 0, confidence: 0.69
1/1 ━━━━━━━━━━━━━━━━━━━━ 2s 2s/step
1/1 ━━━━━━━━━━━━━━━━━━━━ 2s 2s/step
KERAS Predicted class: [1], confidence: [0.002476   0.99752396]
TF SavedModel Predicted class: 0, confidence: 0.72

Keras prediction is right.
ONNX and TF SavedModel are wrong.

[sonali-kumari1]

Hi @lambda-science -

I have tested your code in this gist file with the latest version of keras(3.10.0). I encountered several warnings and errors during onnx export. It appears that tf2onnx does not fully support some data augmentation layers, specifically having issues with StatelessRandomUniformV2 and ImageProjectiveTransformV3. These operations are internally used by layers like RandomBrightness, RandomContrast, RandomFlip, RandomRotation, RandomZoom and RandomTranslation. You can find a list of supported TensorFlow ops and their mapping to onnx here.

Additionally, I have also tested your code with model.keras without onnx conversion here and it seems to works fine. This indicates that the issue is specifically related to tf2onnx. Thanks!