Skip to content

Keras ONNX export not respecting input node name #21637

@ankandrew

Description

@ankandrew

When using torch backend and exporting to onnx, it doesn't seem to respect the name of the InputSpec.

import os
os.environ["KERAS_BACKEND"] = "torch"
os.environ["CUDA_VISIBLE_DEVICES"] = ""

import numpy as np
import keras

inputs = keras.Input(shape=(10, 16))
x = keras.layers.MultiHeadAttention(num_heads=2, key_dim=16)(inputs, inputs)
outputs = keras.layers.Dense(8)(x)
model = keras.Model(inputs=inputs, outputs=outputs)

model.compile(optimizer="adam", loss="mse")

X_dummy = np.random.rand(5, 10, 16).astype(np.float32)
y_dummy = np.random.rand(5, 10, 8).astype(np.float32)

model.fit(X_dummy, y_dummy, epochs=1, verbose=1)

spec = [keras.InputSpec(name="input", shape=(1, 10, 16), dtype="float32")]
model.export("test.onnx", format="onnx", verbose=False, input_signature=spec)

Produces the following names:

Image

Metadata

Metadata

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions