Fix caffe2 relay frontend (#2733)

python/tvm/relay/frontend/caffe2.py

@@ -133,24 +133,14 @@ class Elemwise(Caffe2OpConverter):
     """
     name = ''
     @classmethod
-    def _math_name_picker(cls, suffix):
-
-        def _impl(attr):
-            if attr.get('broadcast', 0):
-                return 'broadcast_' + suffix
-            return 'elemwise_' + suffix
-
-        return _impl
-
-    @classmethod
     def _impl(cls, inputs, args, params):
         assert len(inputs) == 2, "Math op take 2 inputs, {} given".format(
             len(inputs))
-        op_name = cls._math_name_picker(cls.name)(args)
-        axis = int(args.get('axis', 0))
+        op_name = cls.name
         conv_ops = ["conv2d", "conv2d_transpose"]
-        if op_name == 'broadcast_add' and inputs[0].attr('op_name') in conv_ops:
+        if args.get('broadcast', 0) and any(x in str(inputs[0]) for x in conv_ops):
             # TODO(zhreshold): remove hard coded infershape
+            axis = int(args.get('axis', 0))
             inputs[1] = _op.expand_dims(inputs[1], axis=axis, num_newaxis=2)
         return get_relay_op(op_name)(*inputs)
 
@@ -214,7 +204,7 @@ class Conv(Caffe2OpConverter):
                 'order': ('data_layout', ("NCHW"), lambda x: x if isinstance(x, str) else x.decode('UTF-8')),
             },
             excludes=[],
-            ignores=[],
+            ignores=_caffe2_internal_args,
             custom_check=dimension_constraint())(inputs[:2], args, params)
         use_bias = len(inputs) == 3
         if use_bias:
@@ -256,7 +246,7 @@ class NormalizePlanarYUV(Caffe2OpConverter):
         mean = _op.expand_dims(inputs[1], axis=2, num_newaxis=2)
         std = _op.expand_dims(inputs[2], axis=2, num_newaxis=2)
 
-        return _op.broadcast_divide(_op.subtract(inputs[0], mean), std)
+        return _op.divide(_op.subtract(inputs[0], mean), std)
 
 
 class ResizeNearest(Caffe2OpConverter):

tests/python/frontend/caffe2/test_forward.py

@@ -4,7 +4,9 @@ from tvm.contrib import graph_runtime
 from tvm.relay.testing.config import ctx_list
 from tvm import relay
 from model_zoo import c2_squeezenet, c2_resnet50, c2_vgg19
-from caffe2.python import workspace
+from caffe2.python import workspace, core
+from caffe2.proto import caffe2_pb2
+from collections import namedtuple
 
 
 def get_tvm_output(model,
@@ -81,7 +83,135 @@ def test_forward_vgg19():
     verify_caffe2_forward_impl(c2_vgg19, (1, 3, 224, 224), (1, 1000))
 
 
+Model = namedtuple('Model', ['init_net', 'predict_net'])
+
+
+def test_elementwise_add():
+    data_shape = (1, 16, 9, 9)
+    init_net = caffe2_pb2.NetDef()
+    init_net.name = 'test_init_net'
+    init_net.external_output[:] = ['A', 'B']
+    init_net.op.extend([
+        core.CreateOperator(
+            'GivenTensorFill',
+            [],
+            ['A'],
+            shape=data_shape,
+            values=np.random.uniform(size=data_shape).flatten().tolist(),
+        ),
+        core.CreateOperator(
+            'GivenTensorFill',
+            [],
+            ['B'],
+            shape=data_shape,
+            values=np.random.uniform(size=data_shape).flatten().tolist(),
+        ),
+    ])
+
+    predict_net = caffe2_pb2.NetDef()
+    predict_net.name = 'test_predict_net'
+    predict_net.external_input[:] = ['A', 'B']
+    predict_net.external_output[:] = ['C']
+    predict_net.op.extend([
+        core.CreateOperator(
+            'Add',
+            ['A', 'B'],
+            ['C'],
+        )
+    ])
+
+    model = Model(init_net, predict_net)
+    verify_caffe2_forward_impl(model, data_shape, data_shape)
+
+
+def test_elementwise_add_with_broadcast():
+    data_shape = (1, 16, 9, 9)
+    init_net = caffe2_pb2.NetDef()
+    init_net.name = 'test_init_net'
+    init_net.external_output[:] = ['A', 'B']
+    init_net.op.extend([
+        core.CreateOperator(
+            'GivenTensorFill',
+            [],
+            ['A'],
+            shape=data_shape,
+            values=np.random.uniform(size=data_shape).flatten().tolist(),
+        ),
+        core.CreateOperator(
+            'GivenTensorFill',
+            [],
+            ['B'],
+            shape=(1,),
+            values=np.random.uniform(size=1).flatten().tolist(),
+        ),
+    ])
+
+    predict_net = caffe2_pb2.NetDef()
+    predict_net.name = 'test_predict_net'
+    predict_net.external_input[:] = ['A', 'B']
+    predict_net.external_output[:] = ['C']
+    predict_net.op.extend([
+        core.CreateOperator(
+            'Add',
+            ['A', 'B'],
+            ['C'],
+            broadcast=1,
+        )
+    ])
+
+    model = Model(init_net, predict_net)
+    verify_caffe2_forward_impl(model, data_shape, data_shape)
+
+
+def test_normalize_yuv():
+    data_shape = (1, 3, 96, 96)
+    init_net = caffe2_pb2.NetDef()
+    init_net.name = 'test_init_net'
+    init_net.external_output[:] = ['A', 'mean', 'std']
+    init_net.op.extend([
+        core.CreateOperator(
+            'GivenTensorFill',
+            [],
+            ['A'],
+            shape=data_shape,
+            values=np.random.uniform(size=data_shape).flatten().tolist(),
+        ),
+        core.CreateOperator(
+            'GivenTensorFill',
+            [],
+            ['mean'],
+            shape=(1, 3,),
+            values=np.random.uniform(size=3).flatten().tolist(),
+        ),
+        core.CreateOperator(
+            'GivenTensorFill',
+            [],
+            ['std'],
+            shape=(1, 3,),
+            values=np.random.uniform(size=3).flatten().tolist(),
+        ),
+    ])
+
+    predict_net = caffe2_pb2.NetDef()
+    predict_net.name = 'test_predict_net'
+    predict_net.external_input[:] = ['A', 'mean', 'std']
+    predict_net.external_output[:] = ['C']
+    predict_net.op.extend([
+        core.CreateOperator(
+            'NormalizePlanarYUV',
+            ['A', 'mean', 'std'],
+            ['C'],
+        )
+    ])
+
+    model = Model(init_net, predict_net)
+    verify_caffe2_forward_impl(model, data_shape, data_shape)
+
+
 if __name__ == '__main__':
     test_forward_squeezenet1_1()
     test_forward_resnet50()
     test_forward_vgg19()
+    test_elementwise_add()
+    test_elementwise_add_with_broadcast()
+    test_normalize_yuv()