The TF model and ONNX model converted from the TF model have different output shapes, when tf.keras.layers.Conv3DTranspose is used
ktsumura opened this issue · 0 comments
ktsumura commented
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
- 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()
- 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 onnx/tensorflow-onnx#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.