Skip to content

Commit

Permalink
Support mode=instance, spatial for MXNet l2_normalize (apache#7062)
Browse files Browse the repository at this point in the history
  • Loading branch information
Trevor Morris authored and tkonolige committed Jan 11, 2021
1 parent 46ce5f0 commit f176113
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 3 deletions.
13 changes: 10 additions & 3 deletions python/tvm/relay/frontend/mxnet.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)


Expand Down
6 changes: 6 additions & 0 deletions tests/python/frontend/mxnet/test_forward.py
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down

0 comments on commit f176113

Please sign in to comment.