From 16b2e0d43e97ac6df8c01201c379d943655c4e74 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Mon, 23 Aug 2021 21:10:17 +0900 Subject: [PATCH 1/2] [AMP] Disallow fp16 conversion for summation-like ops --- python/tvm/relay/transform/mixed_precision.py | 15 ++++++------- tests/python/relay/test_to_mixed_precision.py | 22 +++++++++++++++++++ 2 files changed, 29 insertions(+), 8 deletions(-) diff --git a/python/tvm/relay/transform/mixed_precision.py b/python/tvm/relay/transform/mixed_precision.py index 1657f895dcd7..fb4d3fa208a8 100644 --- a/python/tvm/relay/transform/mixed_precision.py +++ b/python/tvm/relay/transform/mixed_precision.py @@ -81,8 +81,6 @@ "divide", "nn.bias_add", "nn.batch_norm", - "sum", - "mean", "sqrt", "shape_of", # Simple activations @@ -107,15 +105,9 @@ # "nn.global_max_pool1d", # does not exist yet "nn.global_max_pool2d", # "nn.global_max_pool3d", # does not exist yet - # "nn.global_avg_pool1d", # does not exist yet - "nn.global_avg_pool2d", - # "nn.global_avg_pool3d", # does not exist yet "nn.adaptive_max_pool1d", "nn.adaptive_max_pool2d", "nn.adaptive_max_pool3d", - "nn.adaptive_avg_pool1d", - "nn.adaptive_avg_pool2d", - "nn.adaptive_avg_pool3d", ] DEFAULT_NEVER_LIST = [ # In general if |f(x)| >> |x| for expected inputs then put the op here. @@ -131,6 +123,13 @@ # Do not allow arange arguments (begin/end) to be fp16. "end" can be a big fp32 number # not representable in fp16. "arange", + # Ops that could involve a large summation are not allowed in fp16. + "nn.global_avg_pool2d", + "nn.adaptive_avg_pool1d", + "nn.adaptive_avg_pool2d", + "nn.adaptive_avg_pool3d", + "sum", + "mean", ] diff --git a/tests/python/relay/test_to_mixed_precision.py b/tests/python/relay/test_to_mixed_precision.py index 99078b7371ba..162fa87288ae 100644 --- a/tests/python/relay/test_to_mixed_precision.py +++ b/tests/python/relay/test_to_mixed_precision.py @@ -240,6 +240,28 @@ def test_do_not_convert_arange(): assert tvm.ir.structural_equal(mod, output_mod) +def test_do_not_convert_summation(): + """Ops that could involve a large summation are not allowed in fp16.""" + shape = [1, 3, 16, 16] + a = relay.var("a", shape=shape) + ops = [ + relay.sum, + relay.mean, + relay.nn.global_avg_pool2d, + lambda inp: relay.nn.adaptive_avg_pool2d(inp, (1, 1)), + ] + for op in ops: + mod = tvm.IRModule.from_expr(op(a)) + mod = tvm.relay.transform.InferType()(mod) + + mod_params = { + "a": np.random.uniform(-1, 1, size=shape).astype("float32"), + } + + output_mod = verify_mixed_precision_output_close(mod, mod_params, atol=0.0, rtol=0) + assert tvm.ir.structural_equal(mod, output_mod) + + def test_green_gray_propagates_simple(): """Conv is a green listed operation, while addition is gray. From 8d790a111fc22f635e652720c93baf4311446ab1 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Tue, 24 Aug 2021 18:44:22 +0900 Subject: [PATCH 2/2] test only structural equality --- tests/python/relay/test_to_mixed_precision.py | 27 +++++++------------ 1 file changed, 9 insertions(+), 18 deletions(-) diff --git a/tests/python/relay/test_to_mixed_precision.py b/tests/python/relay/test_to_mixed_precision.py index 162fa87288ae..472f98715ec5 100644 --- a/tests/python/relay/test_to_mixed_precision.py +++ b/tests/python/relay/test_to_mixed_precision.py @@ -221,12 +221,9 @@ def test_do_not_convert_softmax(): b = relay.nn.softmax(a) mod = tvm.IRModule.from_expr(b) mod = tvm.relay.transform.InferType()(mod) - - mod_params = { - "a": np.random.uniform(-1, 1, size=shape).astype("float32"), - } - output_mod = verify_mixed_precision_output_close(mod, mod_params, atol=0.0, rtol=0) - assert tvm.ir.structural_equal(mod, output_mod) + out_mod = ToMixedPrecision("float16")(mod) + orig_mod = tvm.relay.transform.InferType()(mod) + assert tvm.ir.structural_equal(orig_mod, out_mod) def test_do_not_convert_arange(): @@ -234,10 +231,9 @@ def test_do_not_convert_arange(): dtype = "float32" arange = relay.arange(relay.const(1, dtype), relay.const(128, dtype)) mod = tvm.IRModule.from_expr(arange) - mod = tvm.relay.transform.InferType()(mod) - - output_mod = verify_mixed_precision_output_close(mod, {}, atol=0.0, rtol=0) - assert tvm.ir.structural_equal(mod, output_mod) + out_mod = ToMixedPrecision("float16")(mod) + orig_mod = tvm.relay.transform.InferType()(mod) + assert tvm.ir.structural_equal(orig_mod, out_mod) def test_do_not_convert_summation(): @@ -252,14 +248,9 @@ def test_do_not_convert_summation(): ] for op in ops: mod = tvm.IRModule.from_expr(op(a)) - mod = tvm.relay.transform.InferType()(mod) - - mod_params = { - "a": np.random.uniform(-1, 1, size=shape).astype("float32"), - } - - output_mod = verify_mixed_precision_output_close(mod, mod_params, atol=0.0, rtol=0) - assert tvm.ir.structural_equal(mod, output_mod) + out_mod = ToMixedPrecision("float16")(mod) + orig_mod = tvm.relay.transform.InferType()(mod) + assert tvm.ir.structural_equal(orig_mod, out_mod) def test_green_gray_propagates_simple():