Update ONNX-GraphSurgeon to v0.3.9

Signed-off-by: Rajeev Rao <rajeevrao@nvidia.com>

Author: rajeevsrao

tools/onnx-graphsurgeon/CHANGELOG.md

@@ -3,6 +3,20 @@
 Dates are in YYYY-MM-DD format.
 
 
+## v0.3.9 (2021-04-20)
+### Changed
+- `fold_constants()` will no longer store values for foldable tensors whose outputs are all foldable.
+    For example, while folding a constant subgraph like `A (constant) -> B -> C`, previously, `B` values
+    would be computed in addition to `C`. With these changes, only `C` values are computed and stored.
+    This can reduce memory usage significantly.
+
+
+## v0.3.8 (2021-04-15)
+### Fixed
+- Fixed a bug where `copy()` would not work with subgraphs that included tensors with the same
+    names as outer graph tensors unless a `tensor_map` was provided.
+
+
 ## v0.3.7 (2021-03-31)
 ### Added
 - `fold_constants()` can now fold `Shape -> Gather` patterns even when the entire shape may not be known.

tools/onnx-graphsurgeon/onnx_graphsurgeon/__init__.py

@@ -5,4 +5,4 @@
 from onnx_graphsurgeon.ir.tensor import Constant, Tensor, Variable
 from onnx_graphsurgeon.util.exception import OnnxGraphSurgeonException
 
-__version__ = "0.3.7"
+__version__ = "0.3.9"

tools/onnx-graphsurgeon/onnx_graphsurgeon/ir/graph.py

@@ -169,6 +169,7 @@ def _get_node_id(self, node):
     def _local_tensors(self):
         local_tensors = {t.name: t for node in self.nodes for t in node.outputs if not t.is_empty()}
         local_tensors.update({t.name: t for t in self.inputs})
+        local_tensors.update({t.name: t for t in self.tensors().values() if isinstance(t, Constant)})
         return local_tensors
 
 
@@ -279,7 +280,7 @@ def cleanup_subgraphs():
         if recurse_subgraphs:
             cleanup_subgraphs()
 
-        G_LOGGER.debug("Cleaning up {:}".format(self.name))
+        G_LOGGER.verbose("Cleaning up {:}".format(self.name))
 
         with self.node_ids():
             # Graph input producers must be removed first so used_node_ids is correct.
@@ -658,7 +659,17 @@ def get_out_node_ids():
 
 
         # Next, evaluate the foldable variables with ONNX-Runtime
-        graph_clone.outputs = [t for t in graph_constants.values() if not isinstance(t, Constant)]
+
+        # Only evaluate foldable values that have non-foldable outputs or are graph outputs.
+        # Otherwise, if all the outputs are foldable, then we can just evaluate the outputs directly.
+        def should_eval_foldable(tensor):
+            non_const = not isinstance(tensor, Constant)
+            is_graph_output = not tensor.outputs
+            has_non_foldable_outputs = any(out.name not in graph_constants for out in tensor.outputs)
+            return non_const and (is_graph_output or has_non_foldable_outputs)
+
+        graph_clone.outputs = [t for t in graph_constants.values() if should_eval_foldable(t)]
+        G_LOGGER.debug("Folding tensors: {:}".format(graph_clone.outputs))
         graph_clone.cleanup(remove_unused_graph_inputs=True)
 
         # Using ._values avoids a deep copy of the values.
@@ -789,9 +800,15 @@ def copy(self, tensor_map: "OrderedDict[str, Tensor]"=None):
         # First, reconstruct each tensor in the graph, but with no inputs or outputs
         tensor_map = copy.copy(misc.default_value(tensor_map, {}))
 
-        local_tensors = self.tensors()
-        local_tensor_copies = {name: tensor.copy() for name, tensor in local_tensors.items()}
+        local_tensor_copies = {}
+        # When we're cloning a subgraph by itself, we need to use `tensors()` to get all
+        # required tensors - even those produced by outer graphs.
+        local_tensor_copies.update({n: t.copy() for n, t in self.tensors().items()})
+        # However, we should prioritize copies already made by the outer graph.
         local_tensor_copies.update(tensor_map)
+        # And locally produced tensors should take precedence over everything else.
+        local_tensor_copies.update({n: t.copy() for n, t in self._local_tensors().items()})
+
 
         def get_tensor(name):
             if not name:

tools/onnx-graphsurgeon/tests/ir/test_graph.py

@@ -612,6 +612,38 @@ def test_copy_with_subgraph(self, nested_graph):
         assert len(subgraph.nodes) == 2
 
 
+    # If the subgraph has a tensor with the same name as the outer graph,
+    # the subgraph copy should include a copy of the subgraph tensor, not the outer
+    # graph tensor.
+    def test_copy_with_subgraph_dup_tensors(self):
+        inp = Variable("input", dtype=np.float32, shape=(4, 5))
+        graph = Graph(inputs=[inp])
+
+        # We'll use shape to distinguish inner/outer tensor
+        subgraph_inp = Variable("input", dtype=np.float32, shape=(1, 2))
+        subgraph = Graph(inputs=[subgraph_inp])
+
+        graph.outputs = [graph.nested(inp, subgraph)]
+
+        graph_copy = graph.copy()
+        assert graph_copy.nodes[0].attrs["body"].inputs[0].shape == (1, 2)
+
+
+    def test_copy_with_subgraph_dup_const_tensors(self):
+        inp = Constant("input", values=np.ones(dtype=np.float32, shape=(4, 5)))
+        graph = Graph()
+
+        # We'll use shape to distinguish inner/outer tensor
+        subgraph_inp = Constant("input", values=np.ones(dtype=np.float32, shape=(1, 2)))
+        subgraph = Graph()
+        subgraph.outputs = [subgraph.identity(subgraph_inp)]
+
+        graph.outputs = [graph.nested(inp, subgraph)]
+
+        graph_copy = graph.copy()
+        assert graph_copy.nodes[0].attrs["body"].nodes[0].inputs[0].shape == (1, 2)
+
+
 @pytest.fixture
 def simple_foldable():
     # Graph:
@@ -783,8 +815,8 @@ def test_shape_of_constant_node(self):
 
     # Cannot fold shape nodes if they have dynamically shaped inputs.
     def test_shape_of_variable_tensor_dynamic_shape(self):
-        graph = Graph()
         var = Variable("var", dtype=np.float32, shape=("", -1, 0, 4))
+        graph = Graph(inputs=[var])
         graph.outputs = [graph.shape(var)]
 
         graph.fold_constants().cleanup()
@@ -795,8 +827,8 @@ def test_shape_of_variable_tensor_dynamic_shape(self):
 
 
     def test_shape_of_variable_tensor_static_shape(self):
-        graph = Graph()
         var = Variable("var", dtype=np.float32, shape=(1, 3, 4))
+        graph = Graph(inputs=[var])
         graph.inputs = [var]
         graph.outputs = [graph.shape(var)]