Conversion of tranpose into reshape goes wrong after 3D NDHWC convolution (FusedBatchNormV3)

[qrsforever]

when i run the model (CNN3d + BN) and convert to onnx, the FusedBatchNormV3 node's perm attr is wrong.

Reproduce

import tensorflow as tf
import tf2onnx
layers = tf.keras.layers
regularizers = tf.keras.regularizers

class ModelCNN3D(tf.keras.models.Model):
    def __init__(self):
        super().__init__(name='ModelCNN3D')
        self.temporal_conv_layer = layers.Conv3D(
            3, 3, padding='same', name='temporal_conv_layer',
            dilation_rate=(1, 1, 1),
            kernel_regularizer=regularizers.l2(1e-6),
            kernel_initializer='he_normal'
        )
        self.temporal_bn_layers = layers.BatchNormalization(name='temporal_bn_layers')                  
        
    def call(self, x):
        x = self.temporal_conv_layer(x)
        x = self.temporal_bn_layers(x)
        x = tf.nn.relu(x)
        return x
    
testmodel = ModelCNN3D()
inputs = np.random.randn(1, 10, 28, 28, 3)
outputs = testmodel(inputs)
tf2onnx.convert.from_keras(
    testmodel,
    input_signature=(tf.TensorSpec(shape=inputs.shape, name="x"),),
    opset=11, output_path=f'{SAVED_MODEL_ROOT}/testmodel.onnx')[0];

Netron (Transpose perm attr):

Logs:

2021-10-20 17:57:21,970 - INFO - tf2onnx.tfonnx: Using tensorflow=2.6.0, onnx=1.10.1, tf2onnx=1.10.0/4245d8
2021-10-20 17:57:21,972 - INFO - tf2onnx.tfonnx: Using opset <onnx, 11>
2021-10-20 17:57:21,980 - INFO - tf2onnx.tf_utils: Computed 0 values for constant folding
2021-10-20 17:57:21,989 - DEBUG - tf2onnx.graph: Making node: Name=Identity, OP=Identity
2021-10-20 17:57:21,996 - DEBUG - tf2onnx.graph: Made node: Identity

blalalala

2021-10-20 17:57:22,072 - INFO - tf2onnx.optimizer: Optimizing ONNX model
2021-10-20 17:57:22,073 - VERBOSE - tf2onnx.optimizer: Apply optimize_transpose
2021-10-20 17:57:22,077 - DEBUG - tf2onnx.optimizer.TransposeOptimizer: cannot move transpose down since it met output node Identity_graph_outputs_Identity__124
2021-10-20 17:57:22,083 - DEBUG - tf2onnx.optimizer.TransposeOptimizer: finish after 3 iteration(s)
2021-10-20 17:57:22,086 - VERBOSE - tf2onnx.optimizer.TransposeOptimizer: Identity -2 (3->1)

full.log

[qrsforever]

Maybe the BatchNorm Op is not support 3D:

tensorflow-onnx/tf2onnx/onnx_opset/nn.py

Lines 975 to 977 in 4245d8d

	conv_convert_inputs(ctx, node, with_kernel=False)

I add spatial=3 in the functon conv_convert_inputs as below:

        spatial = 3 if len(ctx.get_shape(node.input[0])) == 5 else 2
        conv_convert_inputs(ctx, node, with_kernel=False, spatial=spatial)

Result:

[hwangdeyu]

Maybe the BatchNorm Op is not support 3D:

tensorflow-onnx/tf2onnx/onnx_opset/nn.py

Lines 975 to 977 in 4245d8d

	conv_convert_inputs(ctx, node, with_kernel=False)

I add spatial=3 in the functon conv_convert_inputs as below:

        spatial = 3 if len(ctx.get_shape(node.input[0])) == 5 else 2
        conv_convert_inputs(ctx, node, with_kernel=False, spatial=spatial)

Result:

Thanks @qrsforever, seems this issue is fixed. closed it.