Merge pull request #1074 from NikolasMarkou/upsample_optimizer

Feature: Added optimization step to remove upsample layers with all ones in scale

Author: guschmue

tests/test_optimizers.py

@@ -1177,6 +1177,64 @@ def test_cast_back_to_back_non_const_mixed_types(self):
         self.run_and_compare(["res", "res2", "res3"], {"u": np.random.randn(1, 2, 3).astype(np.float32)}, model_proto,
                              "Cast", 5)
 
+    @check_opset_max_version(8, "until opset 8 scales is in attributes")
+    def test_upsample_all_ones_removed(self):
+        shape = (1, 1, 32, 32)
+        node1 = helper.make_node(
+            op_type="Upsample",
+            inputs=["X"],
+            outputs=["Y"],
+            scales=[1., 1., 1., 1.],
+            name="upsample1")
+
+        graph = helper.make_graph(
+            [node1],
+            "test_upsample_all_ones",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, shape)],
+            [helper.make_tensor_value_info("Y", TensorProto.FLOAT, shape)],
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+
+        self.run_and_compare(
+            ["Y"],
+            {"X": np.random.randn(*shape).astype(np.float32)},
+            model_proto,
+            "Upsample",
+            0)
+
+    @check_opset_min_version(9, ">= 9 scales is in input[1]")
+    @check_opset_max_version(9, "Upscale is deprecated in opsets >= 10")
+    def test_upsample_all_ones_removed_in_input(self):
+        shape = (1, 1, 32, 32)
+        const_tensor = helper.make_tensor(
+            name="S",
+            data_type=TensorProto.FLOAT,
+            dims=(1, 4),
+            vals=np.array([1.0, 1.0, 1.0, 1.0], dtype=np.float32))
+        node0 = helper.make_node("Constant", [], ["S"], value=const_tensor)
+        node1 = helper.make_node(
+            op_type="Upsample",
+            inputs=["X", "S"],
+            outputs=["Y"],
+            name="upsample1")
+
+        graph = helper.make_graph(
+            [node0, node1],
+            "test_upsample_all_ones",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, shape)],
+            [helper.make_tensor_value_info("Y", TensorProto.FLOAT, shape)],
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+
+        self.run_and_compare(
+            ["Y"],
+            {"X": np.random.randn(*shape).astype(np.float32)},
+            model_proto,
+            "Upsample",
+            0)
+
 
 if __name__ == "__main__":
     unittest_main()

tf2onnx/optimizer/__init__.py

@@ -15,11 +15,13 @@
 from .transpose_optimizer import TransposeOptimizer
 from .loop_optimizer import LoopOptimizer
 from .back_to_back_optimizer import BackToBackOptimizer
+from .upsample_optimizer import UpsampleOptimizer
 from .. import logging
 
 # optimizer sequence need to be considered carefully
 _optimizers = OrderedDict([
     ("optimize_transpose", TransposeOptimizer),
+    ("remove_redundant_upsample", UpsampleOptimizer),
     ("fold_constants", ConstFoldOptimizer),
     ("loop_optimizer", LoopOptimizer),
     # merge_duplication should be used after optimize_transpose

tf2onnx/optimizer/upsample_optimizer.py

@@ -0,0 +1,51 @@
+"""Resize Optimizer.
+    Replace resize operations with all ones in scale with Identity nodes
+"""
+
+from __future__ import unicode_literals
+
+import numpy as np
+
+from .optimizer_base import GraphOptimizerBase
+
+# pylint: disable=logging-not-lazy,unused-argument,missing-docstring,unused-variable,arguments-differ
+
+
+class UpsampleOptimizer(GraphOptimizerBase):
+    """Upsample Optimizer."""
+
+    def __init__(self):  # pylint: disable=useless-super-delegation
+        super(UpsampleOptimizer, self).__init__()
+        self._g = None
+
+    def _optimize(self, graph):
+        return self._apply_optimization(
+            graph,
+            self._optimize_at_current_graph_level)
+
+    def _optimize_at_current_graph_level(self, graph):
+        self._g = graph
+        # replace upsample node with all ones in scale with identity node
+        for n in self._g.get_nodes():
+            if n.type == "Upsample":
+                node_changed = False
+                # upsample in opset <=8 has scales in attributes
+                if self._g.opset <= 8:
+                    scales = n.get_attr_value("scales")
+                    if scales and all([float(s) == 1. for s in scales]):
+                        n.type = "Identity"
+                        node_changed = True
+                # upsample in opset >= 9 has scales in input[1]
+                if self._g.opset >= 9 and len(n.input) == 2:
+                    scales_input = n.inputs[1]
+
+                    if scales_input.is_const() and \
+                            np.all(scales_input.get_tensor_value(as_list=False) == 1.):
+                        n.type = "Identity"
+                        n.input = [n.input[0]]
+                        node_changed = True
+                if node_changed:
+                    self.logger.debug("replacing " + n.name +
+                                      " with Identity operation ")
+
+        return self._g