[NNVM][ONNX] Squeeze and Unsqueese operators. (#1339)

nnvm/python/nnvm/frontend/onnx.py

@@ -436,6 +436,16 @@ class Cast(OnnxOpConverter):
         return AttrCvt(op_name='cast', transforms={'to': 'dtype'})(inputs, attr)
 
 
+class Unsqueeze(OnnxOpConverter):
+    """ Operator converter for Unsqueeze.
+    """
+
+    @classmethod
+    def _impl_v1(cls, inputs, attr, params):
+        for axes in attr['axes']:
+            inputs[0] = _sym.expand_dims(inputs[0], axis=axes, num_newaxis=1)
+        return inputs[0]
+
 # compatible operators that do NOT require any conversion.
 _identity_list = []
 
@@ -543,7 +553,8 @@ def _get_convert_map(opset):
         # 'Slice'
         'Transpose': AttrCvt('transpose', {'perm': 'axes'}),
         # 'Gather'
-        # 'Squeeze'
+        'Squeeze': Renamer('squeeze'),
+        'Unsqueeze': Unsqueeze.get_converter(opset),
         'Pad': Pad.get_converter(opset),
         'Shape': Shape.get_converter(opset),
     }

nnvm/tests/python/frontend/onnx/test_forward.py

@@ -107,6 +107,43 @@ def test_reshape_like():
 
     np.testing.assert_allclose(ref_shape, tvm_out.shape)
 
+def test_squeeze():
+    in_shape = (1, 3, 1, 3, 1, 1)
+    out_shape = (3, 3)
+    y = helper.make_node("Squeeze", ['in'], ['out'])
+
+    graph = helper.make_graph([y],
+                              'squeeze_test',
+                              inputs = [helper.make_tensor_value_info("in", TensorProto.FLOAT, list(in_shape))],
+                              outputs = [helper.make_tensor_value_info("out", TensorProto.FLOAT, list(out_shape))])
+
+    model = helper.make_model(graph, producer_name='squeeze_test')
+
+    for target, ctx in ctx_list():
+        x = np.random.uniform(size=in_shape)
+        tvm_out = get_tvm_output(model, x, target, ctx, out_shape, 'float32')
+
+    np.testing.assert_allclose(out_shape, tvm_out.shape)
+
+def test_unsqueeze():
+    in_shape = (3, 3)
+    axis = (0, 3, 4)
+    out_shape = (1, 3, 3, 1, 1)
+    y = helper.make_node("Unsqueeze", ['in'], ['out'], axes=list(axis))
+
+    graph = helper.make_graph([y],
+                              'squeeze_test',
+                              inputs = [helper.make_tensor_value_info("in", TensorProto.FLOAT, list(in_shape))],
+                              outputs = [helper.make_tensor_value_info("out", TensorProto.FLOAT, list(out_shape))])
+
+    model = helper.make_model(graph, producer_name='squeeze_test')
+
+    for target, ctx in ctx_list():
+        x = np.random.uniform(size=in_shape)
+        tvm_out = get_tvm_output(model, x, target, ctx, out_shape, 'float32')
+
+    np.testing.assert_allclose(out_shape, tvm_out.shape)
+
 if __name__ == '__main__':
     # verify_super_resolution_example()
     # verify_squeezenet1_1()
@@ -114,3 +151,5 @@ if __name__ == '__main__':
     verify_resnet18()
     test_reshape()
     test_reshape_like()
+    test_squeeze()
+    test_unsqueeze()