move mkconst to graph

Author: RandyShuai

tf2onnx/graph.py

@@ -405,6 +405,7 @@ def __init__(self, nodes, output_shapes=None, dtypes=None, target=None, opset=No
         if target is None:
             target = []
         self._nodes = []
+        self._consts = {}
         self._nodes_by_name = {}
         self._output_to_node_name = {}
         self.shapes = {}
@@ -484,6 +485,14 @@ def inputs(self):
                 all_inputs.append(n)
         return all_inputs
 
+    def make_consts(self, values, np_type=np.int64, skip_conversion=False, raw=True):
+        """create list of consts of same type"""
+        consts = []
+        for value in values:
+            np_val = np.array(value).astype(np_type)
+            consts.append(self.make_const(utils.make_name("const"), np_val, skip_conversion, raw).output[0])
+        return consts
+
     def make_const(self, name, np_val, skip_conversion=False, raw=True):
         """Make a new constant in the graph.
         Args:
@@ -492,6 +501,11 @@ def make_const(self, name, np_val, skip_conversion=False, raw=True):
             skip_conversion: bool, indicate whether this created node would be mapped during conversion.
             raw: whether to store data at field of raw_data or the specific field according to its dtype
         """
+
+        key = str(np_val) + "_" + str(np_val.dtype)
+        if key in self._consts:
+            return self._consts[key]
+
         if raw:
             onnx_tensor = numpy_helper.from_array(np_val, name)
         else:
@@ -500,6 +514,8 @@ def make_const(self, name, np_val, skip_conversion=False, raw=True):
         dtype = onnx_tensor.data_type
         node = self.make_node("Const", [], outputs=[name], name=name, attr={"value": onnx_tensor},
                               skip_conversion=skip_conversion, dtypes=[dtype], infer_shape_dtype=False)
+
+        self._consts[key] = node
         self.set_shape(name, np_val.shape)
         self.set_dtype(name, utils.map_numpy_to_onnx_dtype(np_val.dtype))
         return node

tf2onnx/onnx_opset/tensor.py

@@ -2223,11 +2223,7 @@ def version_12(cls, ctx, node, **kwargs):
         # Assemble MatrixDiagV3 by ReverseSequence
         argc = len(node.input)
 
-        def mkconsts(values):
-            return [ctx.make_const(utils.make_name('const'), \
-                                   np.array(value).astype(np.int64)).output[0] for value in values]
-
-        minus_two, minus_one, zeo, one, two = mkconsts([[-2], [-1], [0], [1], [2]])
+        minus_two, minus_one, zeo, one, two = ctx.make_consts([[-2], [-1], [0], [1], [2]])
 
         def mknode(op, args, **kwargs):
             return ctx.make_node(op, args, **kwargs).output[0]
@@ -2554,11 +2550,8 @@ class MatrixSetDiagV3:
     @classmethod
     def version_12(cls, ctx, node, **kwargs):
         # Assemble MatrixSetDiagV3 by MatrixDiagPartV3 and MatrixDiagV3
-        def mkconsts(values):
-            return [ctx.make_const(utils.make_name('const'), \
-                                   np.array(value).astype(np.int64)).output[0] for value in values]
 
-        minus_two, minus_one, zeo, one = mkconsts([[-2], [-1], [0], [1]])
+        minus_two, minus_one, zeo, one = ctx.make_consts([[-2], [-1], [0], [1]])
 
         def mknode(op, args, **kwargs):
             return ctx.make_node(op, args, **kwargs).output[0]