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

Author: NikolasMarkou

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,32 @@
+"""Resize Optimizer.
+    Replace resize operations with all ones in scale with Identity nodes
+"""
+
+from __future__ import unicode_literals
+
+from .optimizer_base import GraphOptimizerBase
+
+# pylint: disable=logging-not-lazy,unused-argument,missing-docstring,unused-variable,arguments-differ
+
+
+class UpsampleOptimizer(GraphOptimizerBase):
+    """Resize Optimizer."""
+
+    def __init__(self):  # pylint: disable=useless-super-delegation
+        super(UpsampleOptimizer, self).__init__()
+
+    def _optimize(self, graph):
+        return self._apply_optimization(
+            graph,
+            self._optimize_at_current_graph_level)
+
+    def _optimize_at_current_graph_level(self, graph):
+        # replace resize operations with all ones in scale with Identity nodes
+        for n in graph.get_nodes():
+            if n.type == "Upsample":
+                scales = n.get_attr_value("scales")
+                if all([s == 1 for s in scales]):
+                    n.type = "Identity"
+                    self.logger.debug("replacing " + n.name +
+                                      " with Identity operation")
+        return graph