Support mode=instance, spatial for MXNet l2_normalize (apache#7062)

tkonolige · Jan 11, 2021 · f176113 · f176113
1 parent 46ce5f0
commit f176113
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/python/tvm/relay/frontend/mxnet.py b/python/tvm/relay/frontend/mxnet.py
@@ -1091,12 +1091,19 @@ def _mx_box_decode(inputs, attrs):
 def _mx_l2_normalize(inputs, attrs):
     new_attrs = {}
     mode = attrs.get_str("mode", "instance")
-    if mode != "channel":
+    if mode == "channel":
+        new_attrs["axis"] = [1]
+    elif mode == "instance":
+        ndim = len(_infer_type(inputs[0]).checked_type.shape)
+        new_attrs["axis"] = list(range(1, ndim))
+    elif mode == "spatial":
+        ndim = len(_infer_type(inputs[0]).checked_type.shape)
+        new_attrs["axis"] = list(range(2, ndim))
+    else:
         raise tvm.error.OpAttributeInvalid(
-            'Value of attribute "mode" must equal "channel" for operator l2_normalize.'
+            'Mode "{}" is not supported for operator l2_normalize.'.format(mode)
         )
     new_attrs["eps"] = attrs.get_float("eps", 1e-10)
-    new_attrs["axis"] = [1]
     return _op.nn.l2_normalize(inputs[0], **new_attrs)
 
 

diff --git a/tests/python/frontend/mxnet/test_forward.py b/tests/python/frontend/mxnet/test_forward.py
@@ -638,6 +638,12 @@ def test_forward_l2_normalize():
     mx_sym = mx.sym.L2Normalization(data, mode="channel")
     verify_mxnet_frontend_impl(mx_sym, (2, 3, 4, 5), (2, 3, 4, 5))
 
+    mx_sym = mx.sym.L2Normalization(data, mode="instance")
+    verify_mxnet_frontend_impl(mx_sym, (2, 3, 4, 5), (2, 3, 4, 5))
+
+    mx_sym = mx.sym.L2Normalization(data, mode="spatial")
+    verify_mxnet_frontend_impl(mx_sym, (2, 3, 4, 5), (2, 3, 4, 5))
+
 
 @tvm.testing.uses_gpu
 def test_forward_logistic_regression_output():