The TF model and ONNX model converted from the TF model have different output shapes, when tf.keras.layers.Conv3DTranspose is used

[ktsumura]

Describe the bug The TF model and ONNX model converted from the TF model have different output shapes, when tf.keras.layers.Conv3DTranspose is used.

To Reproduce

1. Run the following code.

import numpy as np
import onnx
import onnxruntime as ort
import tensorflow as tf
import tf2onnx


def transposed_conv_3d_demo():
    tconv_func = tf.keras.layers.Conv3DTranspose(
        256,
        (1, 4, 4),
        strides=(1, 2, 2),
        padding='same',
        output_padding=None,
        data_format='channels_first',
        dilation_rate=(1, 1, 1),
        activation=None,
        use_bias=True)

    input_image = np.ndarray([1, 1, 5, 8, 8]).astype(np.float32)

    # Create a keras model
    inputs = {'in_image': tf.keras.Input(shape=[1, None, 8, 8], name='in_image')}
    outputs = {'out_image': tconv_func(inputs['in_image'])}
    keras_model = tf.keras.models.Model(inputs=inputs, outputs=outputs)
    keras_pred = keras_model.predict(input_image)
    print('The 3D keras model output shape: {}'.format(str(keras_pred['out_image'].shape)))

    # Create and save an ONNX model
    input_signature = [tf.TensorSpec([None, 1, None, 8, 8], tf.float32, name='in_image')]
    onnx_model, _ = tf2onnx.convert.from_keras(keras_model, input_signature, opset=18)
    onnx.save(onnx_model, ".//onnx_model.onnx")

    # ONNX runtime inference
    session = ort.InferenceSession(".//onnx_model.onnx", providers=['CUDAExecutionProvider'])
    onnx_pred = session.run(None, {"in_image": input_image})
    print('The 3D ONNX model output shape: {}'.format(onnx_pred[0].shape))


def transposed_conv_2d_demo():
    # Conv2DTranspose works (with onnx 1.14.0)
    tconv_func = tf.keras.layers.Conv2DTranspose(
        256,
        (4, 4),
        strides=(2, 2),
        padding='same',
        output_padding=None,
        data_format='channels_first',
        dilation_rate=(1, 1),
        activation=None,
        use_bias=True)

    input_image = np.ndarray([1, 1, 8, 8]).astype(np.float32)

    # Create a keras model
    inputs = {'in_image': tf.keras.Input(shape=[1, 8, 8], name='in_image')}
    outputs = {'out_image': tconv_func(inputs['in_image'])}
    keras_model = tf.keras.models.Model(inputs=inputs, outputs=outputs)
    keras_pred = keras_model.predict(input_image)
    print('The 2D keras model output shape: {}'.format(str(keras_pred['out_image'].shape)))

    # Create and save an ONNX model
    input_signature = [tf.TensorSpec([None, 1, 8, 8], tf.float32, name='in_image')]
    onnx_model, _ = tf2onnx.convert.from_keras(keras_model, input_signature, opset=18)
    onnx.save(onnx_model, ".//onnx_model.onnx")

    # ONNX runtime inference
    session = ort.InferenceSession(".//onnx_model.onnx", providers=['CUDAExecutionProvider'])
    onnx_pred = session.run(None, {"in_image": input_image})
    print('The 2D ONNX model output shape: {}'.format(onnx_pred[0].shape))


if __name__ == '__main__':
    transposed_conv_3d_demo()
    transposed_conv_2d_demo()

2. The output shapes are as follows.

1/1 [==============================] - 2s 2s/step
The 3D keras model output shape: (1, 256, 5, 16, 16)
The 3D ONNX model output shape: (1, 256, 5, 16, 18)
2 0 2 3 - 1 2 - 0 4   1 3 : 0 3 : 2 4 . 6 3 9 0 2 8 9   [ W : o n n x r u n t i m e : ,   s e s s i o n _ s t a t e . c c : 1 1 6 2   o n n x r u n t i m e : : V e r i f y E a c h N o d e I s A s s i g n e d T o A n E p ]   S o m e   n o d e s   w e r e   n o t   a s s i g n e d   t o   t h e   p r e f e r r e d   e x e c u t i o n   p r o v i d e r s   w h i c h   m a y   o r   m a y   n o t   h a v e   a n   n e g a t i v e   i m p a c t   o n   p e r f o r m a n c e .   e . g .   O R T   e x p l i c i t l y   a s s i g n s   s h a p e   r e l a t e d   o p s   t o   C P U   t o   i m p r o v e   p e r f . 
 2 0 2 3 - 1 2 - 0 4   1 3 : 0 3 : 2 4 . 6 3 9 4 3 1 3   [ W : o n n x r u n t i m e : ,   s e s s i o n _ s t a t e . c c : 1 1 6 4   o n n x r u n t i m e : : V e r i f y E a c h N o d e I s A s s i g n e d T o A n E p ]   R e r u n n i n g   w i t h   v e r b o s e   o u t p u t   o n   a   n o n - m i n i m a l   b u i l d   w i l l   s h o w   n o d e   a s s i g n m e n t s . 
 2 0 2 3 - 1 2 - 0 4   1 3 : 0 3 : 2 4 . 6 4 5 1 5 3 5   [ W : o n n x r u n t i m e : ,   e x e c u t i o n _ f r a m e . c c : 8 5 7   o n n x r u n t i m e : : E x e c u t i o n F r a m e : : V e r i f y O u t p u t S i z e s ]   E x p e c t e d   s h a p e   f r o m   m o d e l   o f   { - 1 , 2 5 6 , - 1 , 1 6 , 1 6 }   d o e s   n o t   m a t c h   a c t u a l   s h a p e   o f   { 1 , 2 5 6 , 5 , 1 6 , 1 8 }   f o r   o u t p u t   c o n v 3 d _ t r a n s p o s e 
1/1 [==============================] - 0s 39ms/step
The 2D keras model output shape: (1, 256, 16, 16)
The 2D ONNX model output shape: (1, 256, 16, 16)

ONNX model file N/A

Python, ONNX, ONNX-TF, Tensorflow version OS: Windows 11 Pro Python: 3.9 onnx: 1.14.0 (onnx 1.15.0 doesn't work due to https://github.com/onnx/tensorflow-onnx/issues/2262) tf2onnx: 1.15.1 onnxruntime-gpu: 1.16.3 Tensorflow: 2.10.1

Additional context I tried both providers=['CUDAExecutionProvider'] and providers=['CPUExecutionProvider'], and the results were the same.