Tensorflow Lambda layer gives Squeeze error

[mchonofsky]

Describe the bug I have a tf2 model that uses a lambda layer. When converted to ONNX, I get a strange error involving a Squeeze op.

Urgency This is blocking a model deployment right now, but the deadline is self-imposed. Resolution as quickly as possible would be great.

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 18.04*): Debian GNU/Linux 11 (bullseye) running on Chromebook VM
• TensorFlow Version: 2.13.0
• Python version: 3.9.2
• ONNXRuntime version (if applicable, e.g. 1.11*): 1.15.1

To Reproduce

• Generate and save the model

fd = np.random.rand(100,16)
fl = np.random.randint(0,1,[100,])
D = 0.4375
! rm -rf ./logs/*
def get_m():
    I = Input(shape=(None, 16))
    d1 = Dense(16, kernel_regularizer='l2', activation='tanh')(I)
    dr1 = Dropout(D)(d1)
    d2 = Dense(16, kernel_regularizer='l2', activation='tanh')(dr1)
    dr2 = Dropout(D)(d2)
    d3 = Dense(16, kernel_regularizer='l2', activation='tanh')(dr2)
    dr3_1 = Dropout(D)(d3)
    dr3_2 = Dropout(D)(d3)
    d4_1 = Dense(1, kernel_regularizer='l2', activation='relu')(dr3_1)
    d4_2 =  Dense(1, kernel_regularizer='l2', activation='relu')(dr3_2)
    c = Concatenate()([d4_1, d4_2, I])
    def custom_layer(x):
        return tf.exp(-x[:, 4] / (1e-9 + x[:, 0]) - x[:, 5] / (1e-9 + x[:, 1]))

    hl = tf.keras.layers.Lambda(custom_layer, output_shape=(1,))(c)
    #d5 = Dense(1, activation='sigmoid')(c)
    return Model(inputs=I, outputs=hl)

model = get_m()
model.compile(loss=tf.keras.losses.BinaryFocalCrossentropy(), optimizer=tf.keras.optimizers.Adam(), metrics=[tf.keras.metrics.AUC(), tf.keras.metrics.Precision(), tf.keras.metrics.Recall()])
model.fit(x=fd, y=fl, epochs=100, validation_split=0.1
tf2onnx.convert.from_keras(model, output_file='tf2_neural_model_2deep.onnx')

• Load the model in onnx-node

ort = require('onnxruntime-node')
sess = await ort.InferenceSession.create('tf2_neural_model_2deep.onnx')
rowVectors = [[0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1.]]
inputTensor = new ort.Tensor('float32', rowVectors.flat(), [rowVectors.length, 1, rowVectors[0].length]);
output = await sess.run({'input_38': inputTensor});

• Get error

2023-09-12 09:25:00.950381910 [E:onnxruntime:, sequential_executor.cc:514 ExecuteKernel] Non-zero status code returned while running Squeeze node. Name:'model_46/lambda_41/strided_slice__288' Status Message: /onnxruntime_src/onnxruntime/core/providers/cpu/tensor/squeeze.h:52 static onnxruntime::TensorShapeVector onnxruntime::SqueezeBase::ComputeOutputShape(const onnxruntime::TensorShape&, const TensorShapeVector&) input_shape[i] == 1 was false. Dimension of input 1 must be 1 instead of 0. shape={1,0,18}

Uncaught:
Error: Non-zero status code returned while running Squeeze node. Name:'model_46/lambda_41/strided_slice__288' Status Message: /onnxruntime_src/onnxruntime/core/providers/cpu/tensor/squeeze.h:52 static onnxruntime::TensorShapeVector onnxruntime::SqueezeBase::ComputeOutputShape(const onnxruntime::TensorShape&, const TensorShapeVector&) input_shape[i] == 1 was false. Dimension of input 1 must be 1 instead of 0. shape={1,0,18}

Model attached here: model.tar.gz

Screenshots None

Additional context None

[mchonofsky]

Update -

• avoiding the Concatenate stage, thus, does not remove the issue:

def get_m():
    I = Input(shape=(None, 16))
    d1 = Dense(16, kernel_regularizer='l2', activation='tanh')(I)
    dr1 = Dropout(D)(d1)
    d2 = Dense(16, kernel_regularizer='l2', activation='tanh')(dr1)
    dr2 = Dropout(D)(d2)
    d3 = Dense(16, kernel_regularizer='l2', activation='tanh')(dr2)
    dr3_1 = Dropout(D)(d3)
    dr3_2 = Dropout(D)(d3)
    d4_1 = Dense(1, kernel_regularizer='l2', activation='relu')(dr3_1)
    d4_2 =  Dense(1, kernel_regularizer='l2', activation='relu')(dr3_2)
    return Model(inputs=I, outputs=tf.exp(- I[:, 3] / (1e-9 + d4_1[:, 0]) - I[:, 4] / (1e-9 + d4_2[:,0])))

• using a simple sigmoid output without an output transformation does remove the issue:

def get_m():
    I = Input(shape=(None, 16))
    d1 = Dense(16, kernel_regularizer='l2', activation='tanh')(I)
    dr1 = Dropout(D)(d1)
    d2 = Dense(16, kernel_regularizer='l2', activation='tanh')(dr1)
    dr2 = Dropout(D)(d2)
    d3 = Dense(16, kernel_regularizer='l2', activation='tanh')(dr2)
    dr3_1 = Dropout(D)(d3)
    dr3_2 = Dropout(D)(d3)
    
    d4_1 = Dense(8, kernel_regularizer='l2', activation='relu')(dr3_1)
    d4_2 = Dense(8, kernel_regularizer='l2', activation='relu')(dr3_2)
    dr4_1 = Dropout(D)(d4_1)
    d5_1 = Dense(4, kernel_regularizer='l2', activation='relu')(dr4_1)
    d4_1 = Dense(1, kernel_regularizer='l2', activation='sigmoid')(d5_1)
    return Model(inputs=I, outputs=d4_1)

Hope this helps! Thank you!