Repository: apple/coremltools

Issue:
Double transpose ops when using image inputs on a TF 2 model
## 🐞Describe the bug
 - There are duplicate transpose ops that should be removed.

## System environment (please complete the following information):
 - coremltools version  (e.g., 3.0b5): 3.20
 - OS (e.g., MacOS, Linux): MacOS
 - any other relevant information:
     - e.g. using TF 2

## Additional context

```
import tensorflow as tf
import coremltools
import tempfile
inpt = tf.keras.layers.Input(shape=(357, 257, 3), batch_size=1)
out = tf.keras.layers.Conv2D(
    filters=4,
    kernel_size=3,
)(inpt)
model = tf.keras.Model(inpt, out)
input_name = model.inputs[0].name.split(":")[0]
input_shape = (1, 357, 257, 3)

print(input_name)
temp = tempfile.NamedTemporaryFile(suffix=".h5")
model.save(temp.name, include_optimizer=False, save_format="h5")
mlmodel = coremltools.converters.tensorflow.convert(
    temp.name,
    inputs={input_name: input_shape},
    image_scale=1.0 / 255.0,
    red_bias=-0.5,
    green_bias=-0.5,
    blue_bias=-0.5,
    image_input_names=input_name
)
mlmodel.save('test.mlmodel')
```
![image](https://user-images.githubusercontent.com/1422280/72943701-c1665100-3d44-11ea-949c-90c1c850a489.png)

I believe this happens because the transpose op is added in the coremltools builder and there isn't an optimization pass to remove it



Discussion:

Resolving patch:
diff --git a/coremltools/converters/nnssa/coreml/graph_pass/mlmodel_passes.py b/coremltools/converters/nnssa/coreml/graph_pass/mlmodel_passes.py
--- a/coremltools/converters/nnssa/coreml/graph_pass/mlmodel_passes.py
+++ b/coremltools/converters/nnssa/coreml/graph_pass/mlmodel_passes.py
@@ -236,4 +236,71 @@ def _remove_disconnected_layers_rec(nn_spec):
     out_degree = _get_blob_out_degree(spec)
     nn_spec = _get_nn_spec(spec)
     # Initiate removal from high level Neural Network spec
-    _remove_disconnected_layers_rec(nn_spec)
\ No newline at end of file
+    _remove_disconnected_layers_rec(nn_spec)
+
+def remove_redundant_transposes(spec):
+    """
+    Removes layers from model specification that are back to back transposes
+    that compose to the identity.
+    """
+
+    def _delete_layers(nn_spec, layers_to_delete):
+        """
+        Given a neural network spec and pairs of transposes to remove, rewire
+        the network to bypass those transposes and remove them from the spec.
+        """
+        nn_layers = nn_spec.layers
+        # First pass: rewire layers to bypass those that will be deleted.
+        for _layer_pair in layers_to_delete:
+            for _nn_layer in nn_layers:
+                if _nn_layer in _layer_pair:
+                    # Skip the layers we're going to delete.
+                    continue
+                if _layer_pair[1].name in _nn_layer.input:
+                    # This layer has one of the deleted as an input. Replace it
+                    # with the deleted layer's input.
+                    idx = [i for i,n in enumerate(_nn_layer.input) if n == _layer_pair[1].name][0]
+                    _nn_layer.input[idx] = _layer_pair[0].input[0]
+        # Second pass: delete the layers.
+        for _layer_pair in layers_to_delete:
+            nn_layers.remove(_layer_pair[0])
+            nn_layers.remove(_layer_pair[1])
+
+    def _find_redundant_transposes(nn_spec):
+        """
+        Search the neural network spec for pairs of transposes that together
+        are the identity, and return a list of those pairs.
+        """
+        nn_layers = nn_spec.layers
+        layers_to_delete = []
+        # This holds the axes definition if the previous layer was a transpose,
+        # otherwise it is None.
+        previous_transpose = None
+        for _layer in nn_layers:
+            layer_type = _layer.WhichOneof('layer')
+            if layer_type != 'transpose' or len(_layer.output) != 1:
+                previous_transpose = None
+                continue
+
+            if not previous_transpose:
+                previous_transpose = {'layer': _layer, 'axes':_layer.transpose.axes}
+            else:
+                # This layer and the previous are transposes. Check if they're each
+                # other's inverses.
+                this_transpose = _layer.transpose.axes
+                composed = [previous_transpose['axes'][i] for i in this_transpose]
+                if all([ax == i for i, ax in enumerate(composed)]):
+                    # These transpose ops are redundant, remove them.
+                    layers_to_delete.append((previous_transpose['layer'], _layer))
+                else:
+                    # Compare this transpose against the next layer.
+                    # TODO: Should we try to combine a sequence if transposes
+                    # into one op?
+                    previous_transpose = {'layer': _layer, 'axes':_layer.transpose.axes}
+        return layers_to_delete
+
+    nn_spec = _get_nn_spec(spec)
+    layers_to_delete = _find_redundant_transposes(nn_spec)
+    _delete_layers(nn_spec, layers_to_delete)
+    if len(layers_to_delete) > 0:
+        print('{} transpose pairs deleted'.format(len(layers_to_delete)))
diff --git a/coremltools/converters/nnssa/coreml/ssa_converter.py b/coremltools/converters/nnssa/coreml/ssa_converter.py
--- a/coremltools/converters/nnssa/coreml/ssa_converter.py
+++ b/coremltools/converters/nnssa/coreml/ssa_converter.py
@@ -219,7 +219,9 @@ def ssa_convert(ssa,
     mlmodel_spec.description.output.extend(modified_output_features_list)
 
     # MLModel passes
-    mlmodel_passes = [remove_disconnected_layers]
+    mlmodel_passes = [remove_disconnected_layers,
+                      remove_redundant_transposes,
+                     ]
     for p in mlmodel_passes:
         p(mlmodel_spec)
 
diff --git a/coremltools/models/neural_network/builder.py b/coremltools/models/neural_network/builder.py
--- a/coremltools/models/neural_network/builder.py
+++ b/coremltools/models/neural_network/builder.py
@@ -3316,8 +3316,9 @@ def set_pre_processing_parameters(self, image_input_names=None, is_bgr=False,
                     _, channels, height, width = [array_shape[e] for e in input_indices]
 
                     if image_format == 'NHWC':
-                        # If input format is 'NHWC', then add transpose
-                        # after the input and replace all use of input
+                        # If input format is 'NHWC' for TF model, it will be
+                        # 'NCHW' for CoreML model. Therefore, add transpose to
+                        # NHWC after the input and replace all use of input
                         # with output of transpose
                         axes = [1, 2, 0]
                         if len(array_shape) == 4:


Test patch:
diff --git a/coremltools/converters/tensorflow/test/test_tf_2x.py b/coremltools/converters/tensorflow/test/test_tf_2x.py
--- a/coremltools/converters/tensorflow/test/test_tf_2x.py
+++ b/coremltools/converters/tensorflow/test/test_tf_2x.py
@@ -670,6 +670,34 @@ def test_softplus(self):
                                  inputs={input_name: (1, 28, 28)},
                                  outputs=[output_name], decimal=3)
 
+    def test_redundant_transpose(self):
+        H = 224
+        W = 224
+        C = 3
+        inputs = tf.keras.layers.Input(shape=(H, W, C), batch_size=1)
+        out = tf.keras.layers.Conv2D(
+            filters=4,
+            kernel_size=3,
+        )(inputs)
+        model = tf.keras.Model(inputs, out)
+        input_name = model.inputs[0].name.split(":")[0]
+        input_shape = (1, H, W, C)
+        output_name = model.outputs[0].name.split(':')[0].split('/')[-1]
+
+        model.save(self.model_path, include_optimizer=False, save_format="h5")
+
+        mlmodel = coremltools.converters.tensorflow.convert(
+            self.model_path,
+            inputs={input_name: input_shape},
+            image_input_names=input_name,
+            outputs=[output_name],
+        )
+
+        spec = mlmodel.get_spec()
+        output_names = [layer.name for layer in spec.neuralNetwork.layers]
+        expected_names = [u'model/conv2d/Conv2D', u'Identity']
+        np.testing.assert_array_equal(output_names, expected_names)
+
 
 if __name__ == '__main__':
     np.random.seed(1984)
diff --git a/coremltools/test/neural_network/test_graph_passes.py b/coremltools/test/neural_network/test_graph_passes.py
--- a/coremltools/test/neural_network/test_graph_passes.py
+++ b/coremltools/test/neural_network/test_graph_passes.py
@@ -5,7 +5,7 @@
 from coremltools.models import MLModel
 from coremltools.models.neural_network.printer import print_network_spec
 from coremltools.converters.nnssa.coreml.graph_pass.mlmodel_passes import \
-        remove_disconnected_layers, transform_conv_crop
+        remove_disconnected_layers, transform_conv_crop, remove_redundant_transposes
 import copy
 import pytest
 
@@ -203,6 +203,81 @@ def test_conv_crop_bn_relu_to_conv_bn_relu_crop(self):
         np.testing.assert_equal('activation', spec.layers[2].WhichOneof('layer'))
         np.testing.assert_equal('crop', spec.layers[3].WhichOneof('layer'))
 
+    def test_redundant_transposes(self):
+
+        def _build_and_test_network(input_size, transpose_layers, expected_layers):
+            """
+            Helper function for testing transpose removal.
+
+            Args:
+                input_size: Size of the input network tensor.
+                transpose_layers: Array of transpose axes definitions.
+                expected_layers: Array of indices into transpose_layers indicating
+                    which of the transpose layers should be present after the
+                    graph pass.
+            """
+            input_features = [('data', datatypes.Array(*input_size))]
+            output_features = [('out', None)]
+            builder = neural_network.NeuralNetworkBuilder(input_features, output_features)
+            last_layer = 'data'
+            for idx, axes in enumerate(transpose_layers):
+                name = 't{}'.format(idx)
+                if idx == len(transpose_layers) - 1:
+                    output_name = 'out'
+                else:
+                    output_name = name + '_out'
+                builder.add_transpose(name=name,
+                                      axes=axes,
+                                      input_name=last_layer,
+                                      output_name=output_name)
+                last_layer = output_name
+
+            spec = builder.spec.neuralNetwork
+            # Check the network before the graph pass.
+            for idx in range(len(transpose_layers)):
+                np.testing.assert_equal('transpose', spec.layers[idx].WhichOneof('layer'))
+            # Run the removal pass.
+            remove_redundant_transposes(builder.spec)
+            # Verify only the expected layers remain.
+            np.testing.assert_equal(len(spec.layers), len(expected_layers))
+            for output_layer_idx, input_layer_idx in enumerate(expected_layers):
+                np.testing.assert_equal(
+                    'transpose',
+                    spec.layers[output_layer_idx].WhichOneof('layer')
+                )
+                np.testing.assert_array_equal(
+                    transpose_layers[input_layer_idx],
+                    spec.layers[output_layer_idx].transpose.axes
+                )
+
+        _build_and_test_network(
+            input_size=[1, 10, 10],
+            # These transposes together are the identity.
+            transpose_layers=[[2, 0, 1], [1, 2, 0]],
+            expected_layers=[],
+        )
+
+        _build_and_test_network(
+            input_size=[1, 10, 10],
+            # These transposes are not inverses.
+            transpose_layers=[[2, 0, 1], [2, 0, 1]],
+            expected_layers=[0, 1],
+        )
+
+        _build_and_test_network(
+            input_size=[1, 1, 10, 10, 3],
+            # First two are the identity, then an extra.
+            transpose_layers=[[2, 4, 1, 0, 3], [3, 2, 0, 4, 1], [1, 0, 2, 3, 4]],
+            expected_layers=[2],
+        )
+
+        _build_and_test_network(
+            input_size=[1, 1, 10, 10, 3],
+            # First is okay, next two are the identity.
+            transpose_layers=[[1, 0, 2, 3, 4], [2, 4, 1, 0, 3], [3, 2, 0, 4, 1]],
+            expected_layers=[0],
+        )
+
 
 if __name__ == '__main__':
     RUN_ALL_TESTS = True