[COREML]Unary ops support added in frontend (#6196)

* [COREML]Unary ops support added in frontend

* Used coreml enums

python/tvm/relay/frontend/coreml.py

@@ -14,7 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# pylint: disable=invalid-name, import-self, unused-argument, unused-variable
+# pylint: disable=invalid-name, import-self, unused-argument, unused-variable, no-else-return
 # pylint: disable=inconsistent-return-statements, import-outside-toplevel
 """CoreML frontend."""
 import math
@@ -350,6 +350,33 @@ def _MinLayerParams(op, inexpr, etab):
     return _min
 
 
+def _UnaryFunctionLayerParams(op, inexpr, etab):
+    op_type = op.type
+    if op_type == op.SQRT:
+        return _op.sqrt(inexpr)
+    elif op_type == op.RSQRT:
+        epsilon = _expr.const(op.epsilon)
+        return _op.rsqrt(inexpr + epsilon)
+    elif op_type == op.INVERSE:
+        epsilon = _expr.const(op.epsilon)
+        return _expr.const(1.0) / (inexpr + epsilon)
+    elif op_type == op.POWER:
+        alpha = _expr.const(op.alpha)
+        return _op.power(inexpr, alpha)
+    elif op_type == op.EXP:
+        return _op.exp(inexpr)
+    elif op_type == op.LOG:
+        return _op.log(inexpr)
+    elif op_type == op.ABS:
+        return _op.abs(inexpr)
+    elif op_type == op.THRESHOLD:
+        alpha = _expr.const(op.alpha)
+        return _op.maximum(inexpr, alpha)
+    else:
+        msg = 'Unary Op type value {} is not supported in frontend CoreML.'
+        raise tvm.error.OpAttributeUnImplemented(msg.format(op_type))
+
+
 _convert_map = {
     'NeuralNetworkMeanImage': _NeuralNetworkMeanImage,
     'NeuralNetworkImageScaler': _NeuralNetworkImageScaler,
@@ -372,6 +399,7 @@ def _MinLayerParams(op, inexpr, etab):
     'AverageLayerParams': _AverageLayerParams,
     'MaxLayerParams': _MaxLayerParams,
     'MinLayerParams': _MinLayerParams,
+    'UnaryFunctionLayerParams': _UnaryFunctionLayerParams,
 }
 
 # SAME padding: https://www.tensorflow.org/api_guides/python/nn

tests/python/frontend/coreml/test_forward.py

@@ -337,6 +337,194 @@ def test_forward_min():
     verify_min((1, 3, 20, 20))
     verify_min((20, 20))
 
+
+def verify_unary_sqrt(input_dim):
+    dtype = 'float32'
+
+    a_np = np.random.uniform(size=input_dim).astype(dtype)
+    ref_val = np.sqrt(a_np)
+
+    inputs = [('input', datatypes.Array(*input_dim))]
+    output = [('output', datatypes.Array(*ref_val.shape))]
+    builder = NeuralNetworkBuilder(inputs, output)
+    builder.add_unary(name="sqrt",
+                      input_name='input',
+                      output_name='output',
+                      mode='sqrt')
+
+    model = cm.models.MLModel(builder.spec)
+    for target, ctx in ctx_list():
+        out = run_tvm_graph(model, target, ctx, [a_np],
+                            ['input'], ref_val.shape, dtype)
+        tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)
+
+
+def verify_unary_rsqrt(input_dim, epsilon=0):
+    dtype = 'float32'
+
+    a_np = np.random.uniform(size=input_dim).astype(dtype)
+    ref_val = 1 / np.sqrt(a_np + epsilon)
+
+    inputs = [('input', datatypes.Array(*input_dim))]
+    output = [('output', datatypes.Array(*ref_val.shape))]
+    builder = NeuralNetworkBuilder(inputs, output)
+    builder.add_unary(name="rsqrt",
+                      input_name='input',
+                      output_name='output',
+                      mode='rsqrt',
+                      epsilon=epsilon)
+
+    model = cm.models.MLModel(builder.spec)
+    for target, ctx in ctx_list():
+        out = run_tvm_graph(model, target, ctx, [a_np],
+                            ['input'], ref_val.shape, dtype)
+        tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)
+
+
+def verify_unary_inverse(input_dim, epsilon=0):
+    dtype = 'float32'
+
+    a_np = np.random.uniform(size=input_dim).astype(dtype)
+    ref_val = 1 / (a_np + epsilon)
+
+    inputs = [('input', datatypes.Array(*input_dim))]
+    output = [('output', datatypes.Array(*ref_val.shape))]
+    builder = NeuralNetworkBuilder(inputs, output)
+    builder.add_unary(name="inverse",
+                      input_name='input',
+                      output_name='output',
+                      mode='inverse',
+                      epsilon=epsilon)
+
+    model = cm.models.MLModel(builder.spec)
+    for target, ctx in ctx_list():
+        out = run_tvm_graph(model, target, ctx, [a_np],
+                            ['input'], ref_val.shape, dtype)
+        tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)
+
+
+def verify_unary_power(input_dim, alpha):
+    dtype = 'float32'
+
+    a_np = np.random.uniform(size=input_dim).astype(dtype)
+    ref_val = np.power(a_np, alpha)
+
+    inputs = [('input', datatypes.Array(*input_dim))]
+    output = [('output', datatypes.Array(*ref_val.shape))]
+    builder = NeuralNetworkBuilder(inputs, output)
+    builder.add_unary(name="power",
+                      input_name='input',
+                      output_name='output',
+                      mode='power',
+                      alpha=alpha)
+
+    model = cm.models.MLModel(builder.spec)
+    for target, ctx in ctx_list():
+        out = run_tvm_graph(model, target, ctx, [a_np],
+                            ['input'], ref_val.shape, dtype)
+        tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)
+
+
+def verify_unary_exp(input_dim):
+    dtype = 'float32'
+
+    a_np = np.random.uniform(size=input_dim).astype(dtype)
+    ref_val = np.exp(a_np)
+
+    inputs = [('input', datatypes.Array(*input_dim))]
+    output = [('output', datatypes.Array(*ref_val.shape))]
+    builder = NeuralNetworkBuilder(inputs, output)
+    builder.add_unary(name="exp",
+                      input_name='input',
+                      output_name='output',
+                      mode='exp')
+
+    model = cm.models.MLModel(builder.spec)
+    for target, ctx in ctx_list():
+        out = run_tvm_graph(model, target, ctx, [a_np],
+                            ['input'], ref_val.shape, dtype)
+        tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)
+
+
+def verify_unary_log(input_dim):
+    dtype = 'float32'
+
+    a_np = np.random.uniform(size=input_dim).astype(dtype)
+    ref_val = np.log(a_np)
+
+    inputs = [('input', datatypes.Array(*input_dim))]
+    output = [('output', datatypes.Array(*ref_val.shape))]
+    builder = NeuralNetworkBuilder(inputs, output)
+    builder.add_unary(name="log",
+                      input_name='input',
+                      output_name='output',
+                      mode='log')
+
+    model = cm.models.MLModel(builder.spec)
+    for target, ctx in ctx_list():
+        out = run_tvm_graph(model, target, ctx, [a_np],
+                            ['input'], ref_val.shape, dtype)
+        tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)
+
+
+def verify_unary_abs(input_dim):
+    dtype = 'float32'
+
+    a_np = np.random.uniform(-100.0, 100.0, size=input_dim).astype(dtype)
+    ref_val = np.abs(a_np)
+
+    inputs = [('input', datatypes.Array(*input_dim))]
+    output = [('output', datatypes.Array(*ref_val.shape))]
+    builder = NeuralNetworkBuilder(inputs, output)
+    builder.add_unary(name="abs",
+                      input_name='input',
+                      output_name='output',
+                      mode='abs')
+
+    model = cm.models.MLModel(builder.spec)
+    for target, ctx in ctx_list():
+        out = run_tvm_graph(model, target, ctx, [a_np],
+                            ['input'], ref_val.shape, dtype)
+        tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)
+
+
+def verify_unary_threshold(input_dim, alpha):
+    dtype = 'float32'
+
+    a_np = np.random.uniform(-100.0, 100.0, size=input_dim).astype(dtype)
+    ref_val = np.maximum(a_np, alpha)
+
+    inputs = [('input', datatypes.Array(*input_dim))]
+    output = [('output', datatypes.Array(*ref_val.shape))]
+    builder = NeuralNetworkBuilder(inputs, output)
+    builder.add_unary(name="threshold",
+                      input_name='input',
+                      output_name='output',
+                      mode='threshold',
+                      alpha=alpha)
+
+    model = cm.models.MLModel(builder.spec)
+    for target, ctx in ctx_list():
+        out = run_tvm_graph(model, target, ctx, [a_np],
+                            ['input'], ref_val.shape, dtype)
+        tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)
+
+
+def test_forward_unary():
+    verify_unary_sqrt((1, 3, 20, 20))
+    verify_unary_rsqrt((1, 3, 20, 20))
+    verify_unary_rsqrt((1, 3, 20, 20), epsilon=1e-6)
+    verify_unary_inverse((1, 3, 20, 20))
+    verify_unary_inverse((1, 3, 20, 20), epsilon=1e-6)
+    verify_unary_power((1, 3, 20, 20), alpha=0.5)
+    verify_unary_power((1, 3, 20, 20), alpha=4)
+    verify_unary_exp((1, 3, 20, 20))
+    verify_unary_log((1, 3, 20, 20))
+    verify_unary_abs((1, 3, 20, 20))
+    verify_unary_threshold((1, 3, 20, 20), alpha=-6.0)
+    verify_unary_threshold((1, 3, 20, 20), alpha=5.0)
+
+
 def verify_image_scaler(input_dim, blue_bias=0.0, green_bias=0.0, red_bias=0.0, image_scale=1.0):
     dtype = 'float32'
     a_np = np.random.uniform(size=input_dim).astype(dtype)
@@ -413,6 +601,7 @@ def test_forward_convolution():
     test_forward_average()
     test_forward_max()
     test_forward_min()
+    test_forward_unary()
     test_mobilenet_checkonly()
     test_resnet50_checkonly()
     test_forward_image_scaler()