Added optimization and test for opset>=9

Author: NikolasMarkou

tests/test_optimizers.py

@@ -1203,6 +1203,37 @@ def test_upsample_all_ones_removed(self):
             "Upsample",
             0)
 
+    @check_opset_min_version(9, ">= 9 scales is in input[1]")
+    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/upsample_optimizer.py

@@ -4,6 +4,8 @@
 
 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
@@ -26,12 +28,24 @@ def _optimize_at_current_graph_level(self, 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"
-                        self.logger.debug("replacing " + n.name +
-                                          " with Identity operation ")
-                # upsample in opset > 8 has scales in input[1]
+                        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