Merge pull request #929 from jignparm/jignparm/finetunedbert

Multiple fixes for Bert Model (fine-tuned)

Author: jignparm

tf2onnx/graph.py

@@ -1067,7 +1067,9 @@ def make_onnx_graph_io(self, ids):
             shape = self.get_shape(name)
 
             utils.make_sure(dtype is not None, "missing output dtype for " + name)
-            utils.make_sure(shape is not None, "missing output shape for " + name)
+            # TODO: allow None output shape or not? e.g. shape=(?,)
+            #utils.make_sure(shape is not None, "missing output shape for " + name)
+            if shape is None: logger.warning("missing output shape for %s", name)
 
             v = utils.make_onnx_inputs_outputs(name, dtype, shape)
             tensor_value_infos.append(v)

tf2onnx/onnx_opset/tensor.py

@@ -196,6 +196,7 @@ def version_1(cls, ctx, node, **kwargs):
             shape = ctx.get_shape(node.input[0])
             utils.make_sure(shape is not None, "squeeze input shape cannot be None")
             axis = [i for i, j in enumerate(shape) if j == 1]
+            if not axis: axis = [0]
         node.set_attr("axes", axis)
 
     @classmethod

tf2onnx/tfonnx.py

@@ -54,21 +54,15 @@ def rewrite_constant_fold(g, ops):
         "Sqrt": np.sqrt,
         "Sub": np.subtract,
     }
-    ref_cnt_per_node = {}
-    for idx, op in enumerate(ops):
-        for op_input in op.inputs:
-            if op_input.name not in ref_cnt_per_node:
-                ref_cnt_per_node[op_input.name] = 0
-            ref_cnt_per_node[op_input.name] += 1
 
     # pylint: disable=too-many-nested-blocks
     keep_looking = True
     while keep_looking:
         keep_looking = False
         for idx, op in enumerate(ops):
             func = func_map.get(op.type)
-            if func is None:
-                continue
+            if func is None: continue
+            if set(op.output) & set(g.outputs): continue
             try:
                 inputs = []
                 for node in op.inputs:
@@ -109,18 +103,14 @@ def rewrite_constant_fold(g, ops):
                     old_node_name = op.name
                     logger.debug("create const node [%s] replacing [%s]", new_node_name, old_node_name)
                     ops[idx] = g.make_const(new_node_name, val)
-                    ref_cnt_per_node[new_node_name] = ref_cnt_per_node[old_node_name]
 
                     logger.debug("replace old output [%s] with new output [%s]", old_output_name, new_output_name)
                     # need to re-write the consumers input name to use the const name
                     consumers = g.find_output_consumers(old_output_name)
                     if consumers:
                         for consumer in consumers:
                             g.replace_input(consumer, old_output_name, new_output_name)
-                    for node in op.inputs:
-                        ref_cnt_per_node[node.name] -= 1
-                        if ref_cnt_per_node[node.name] == 0:
-                            g.remove_node(node.name)
+
                     # keep looking until there is nothing we can fold.
                     # We keep the graph in topological order so if we folded,
                     # the result might help a following op.
@@ -459,8 +449,8 @@ def compat_handler(ctx, node, **kwargs):
 
     # pre-processing graph rewrites
     # bi-directional re-writer should be placed after single directional re-writer
-    rewriters = [rewrite_quantize_and_dequantize, rewrite_transpose, rewrite_flatten, rewrite_gemm,
-                 rewrite_random_uniform, rewrite_random_uniform_fold_const,
+    rewriters = [rewrite_constant_fold, rewrite_quantize_and_dequantize, rewrite_transpose, rewrite_flatten,
+                 rewrite_gemm, rewrite_random_uniform, rewrite_random_uniform_fold_const,
                  rewrite_random_normal, rewrite_dropout, rewrite_eye,
                  rewrite_leakyrelu, rewrite_thresholded_relu, rewrite_conv2d_with_pad,
                  rewrite_single_direction_lstm, rewrite_bi_direction_lstm,