Does tf2onnx convert the model output type?

[matwey]

Hello,

I am running tf2onnx 1.16.1, tensorflow 2.12.0, onnxruntime 1.18.0.

I have the following simple tensorflow keras model:

    model = tf.keras.Sequential([
        tf.keras.layers.Dense(HIDDEN_DIM, activation="linear"),
        tf.keras.layers.Dense(OUTPUT_DIM, activation="softmax", name="output"),
    ])

When trained, It predicts a vector of probabilities:

[[0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00
  9.9999994e-01]
 [7.1569558e-25 9.9751562e-01 1.3877662e-04 5.5733459e-35 1.3489545e-18
  2.3455818e-03]
 [9.5909646e-24 1.6532872e-15 6.3794343e-25 0.0000000e+00 8.0798064e-13
  9.9999994e-01]
 [4.5122406e-15 2.0971049e-24 0.0000000e+00 9.9999994e-01 1.7789930e-27
  0.0000000e+00]
 [9.9995691e-01 4.3005555e-05 8.8478760e-32 1.9627071e-34 1.1105422e-10
  6.1429810e-23]
 [9.7836558e-15 9.9609488e-01 3.9010276e-03 4.1101748e-06 1.2712012e-12
  4.3790528e-19]
 [0.0000000e+00 0.0000000e+00 0.0000000e+00 9.9999994e-01 1.5491074e-38
  0.0000000e+00]
 [0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00
  9.9999994e-01]
 [2.7403751e-30 0.0000000e+00 0.0000000e+00 1.0000000e+00 1.5874274e-33
  0.0000000e+00]
 [9.9987257e-01 1.2731102e-04 0.0000000e+00 0.0000000e+00 1.1552715e-07
  0.0000000e+00]]

Then I save the model as the following:

    spec = (tf.TensorSpec((None, INPUT_DIM), tf.float32, name="input"),)
    output_path = "model.onnx"
    model_proto, _ = tf2onnx.convert.from_keras(model, input_signature=spec, output_path=output_path)

But when I am trying to load and reuse the model with the same data the result is inconsistent:

[[0. 0. 0. 1. 0. 0.]
 [0. 0. 0. 0. 1. 0.]
 [0. 0. 0. 0. 0. 1.]
 [0. 0. 0. 0. 0. 1.]
 [0. 0. 0. 1. 0. 0.]
 [0. 0. 0. 0. 0. 1.]
 [0. 0. 0. 0. 0. 1.]
 [0. 0. 0. 0. 0. 1.]
 [0. 0. 0. 0. 0. 1.]
 [0. 0. 0. 0. 0. 1.]]

What do I do wrong?