Merge pull request #371 from nbcsm/expanddims

shape inference support ExpandDims

Author: nbcsm

tf2onnx/shape_inference.py

@@ -175,7 +175,22 @@ def infer_shape_for_node(g, node):
 
         g.set_shape(node.output[0], new_shape)
         log.debug("set %s node [%s] with new shape %s", node.type, node.output[0], new_shape)
+        return True
+
+    if node.type == "ExpandDims":
+        # https://www.tensorflow.org/api_docs/python/tf/expand_dims
+        input_shape = g.get_shape(node.input[0])
+        dim_node = node.inputs[1]
+        if input_shape is None or not dim_node.is_const():
+            return False
 
+        dim = dim_node.get_tensor_value()
+        if dim < 0:
+            dim = dim + len(input_shape) + 1
+
+        new_shape = input_shape[:dim] + [1] + input_shape[dim:]
+        g.set_shape(node.output[0], new_shape)
+        log.debug("set [%s] with new shape %s", node.output[0], new_shape)
         return True
 
     return False

tf2onnx/tfonnx.py

@@ -881,9 +881,10 @@ def expanddims_op(ctx, node, name, args):
     dim_node = node.inputs[1]
     if dim_node.is_const():
         node.type = "Unsqueeze"
-        input_rank = len(ctx.get_shape(node.input[0]))
         dim = dim_node.get_tensor_value()
-        dim = dim + input_rank + 1 if dim < 0 else dim
+        if dim < 0:
+            input_rank = len(ctx.get_shape(node.input[0]))
+            dim = dim + input_rank + 1
         node.set_attr("axes", [dim])
         ctx.remove_input(node, node.input[1])
         return
@@ -907,9 +908,10 @@ def expanddims_op7(ctx, node, name, args):
     dim_node = node.inputs[1]
     if dim_node.is_const():
         node.type = "Unsqueeze"
-        input_rank = len(ctx.get_shape(node.input[0]))
         dim = dim_node.get_tensor_value()
-        dim = dim + input_rank + 1 if dim < 0 else dim
+        if dim < 0:
+            input_rank = len(ctx.get_shape(node.input[0]))
+            dim = dim + input_rank + 1
         node.set_attr("axes", [dim])
         ctx.remove_input(node, node.input[1])
         return