From dc28a4340917fa8ceb082365e9c107b9b3543651 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Thu, 23 Sep 2021 11:12:22 +0800 Subject: [PATCH] fix(mgb/bn): fix empty tensor input problem and other minor problems GitOrigin-RevId: 1708b76cb83e90cb91e34c58d77329a67cd4a792 --- .../test/integration/test_converge_with_gradient_clip.py | 6 +++--- imperative/python/test/integration/test_save_load.py | 2 -- imperative/python/test/unit/core/test_dtype_quant.py | 1 - imperative/python/test/unit/core/test_imperative_rt.py | 4 ---- imperative/python/test/unit/data/test_dataloader.py | 2 +- imperative/python/test/unit/data/test_pre_dataloader.py | 3 +-- imperative/python/test/unit/data/test_transform.py | 5 +++-- imperative/python/test/unit/functional/test_elemwise.py | 6 ++---- imperative/python/test/unit/functional/test_functional.py | 5 ++--- src/opr/impl/dnn/batch_norm.cpp | 8 +++++++- 10 files changed, 19 insertions(+), 23 deletions(-) diff --git a/imperative/python/test/integration/test_converge_with_gradient_clip.py b/imperative/python/test/integration/test_converge_with_gradient_clip.py index 5a91c5981..8ef05d243 100644 --- a/imperative/python/test/integration/test_converge_with_gradient_clip.py +++ b/imperative/python/test/integration/test_converge_with_gradient_clip.py @@ -107,8 +107,9 @@ def test_training_converge(test_traced_module): optim.clip_grad_value(net.parameters(), lower=-0.1, upper=0.1) opt.step() losses.append(loss.numpy()) - print(np.mean(losses[-100:])) - assert np.mean(losses[-100:]) < 0.1, "Final training Loss must be low enough" + assert ( + np.mean(losses[-100:]) < 0.1 + ), "Final training Loss must be low enough, get {}".format(np.mean(losses[-100:])) ngrid = 10 x = np.linspace(-1.0, 1.0, ngrid) @@ -118,7 +119,6 @@ def test_training_converge(test_traced_module): data = mge.tensor(np.concatenate((xx, yy), axis=1).astype(np.float32)) pred = infer(data) precision = calculate_precision(data.numpy(), pred.numpy()) - print("precision=", precision) assert precision == 1.0, "Test precision must be high enough, get {}".format( precision ) diff --git a/imperative/python/test/integration/test_save_load.py b/imperative/python/test/integration/test_save_load.py index f80af3b8c..c1cab14d2 100644 --- a/imperative/python/test/integration/test_save_load.py +++ b/imperative/python/test/integration/test_save_load.py @@ -74,7 +74,6 @@ def test_save_load(): optim.step() model_name = "simple.pkl" - print("save to {}".format(model_name)) mge.save( { @@ -93,7 +92,6 @@ def test_save_load(): net.load_state_dict(checkpoint["state_dict"]) optim = optimizer.SGD(net.parameters(), lr=1.0, momentum=0.9) optim.load_state_dict(checkpoint["opt_state"]) - print("load done") os.remove("simple.pkl") with gm: diff --git a/imperative/python/test/unit/core/test_dtype_quant.py b/imperative/python/test/unit/core/test_dtype_quant.py index 0fbd99e23..2f2c651fe 100644 --- a/imperative/python/test/unit/core/test_dtype_quant.py +++ b/imperative/python/test/unit/core/test_dtype_quant.py @@ -165,7 +165,6 @@ def test_dtype_int4_ffi_handle(): device = "xpux" shape = (3, 3, 3) data = np.random.random(shape).astype(np.float32) * 5 - 1 - print(data) def identity(x): return x diff --git a/imperative/python/test/unit/core/test_imperative_rt.py b/imperative/python/test/unit/core/test_imperative_rt.py index 04b08307f..14522a617 100644 --- a/imperative/python/test/unit/core/test_imperative_rt.py +++ b/imperative/python/test/unit/core/test_imperative_rt.py @@ -25,10 +25,6 @@ def test_basic_interface(): cf.name = "megengine.core" cf.dtype = "float32" cf.comp_node_arr = ["xpux"] - print(cf.name) - print(cf.dtype) - print(cf.comp_node_arr) - print(cf.comp_node) cf.comp_node_arr = ["xpux", "xpux:1"] with pytest.raises(ValueError): cf.comp_node diff --git a/imperative/python/test/unit/data/test_dataloader.py b/imperative/python/test/unit/data/test_dataloader.py index 789fe8b50..0bfc8e495 100644 --- a/imperative/python/test/unit/data/test_dataloader.py +++ b/imperative/python/test/unit/data/test_dataloader.py @@ -203,7 +203,7 @@ def test_dataloader_parallel_worker_exception(): pass def apply(self, input): - y = x + 1 + raise RuntimeError("test raise error") return input dataloader = DataLoader( diff --git a/imperative/python/test/unit/data/test_pre_dataloader.py b/imperative/python/test/unit/data/test_pre_dataloader.py index 75cd197f3..24394a41b 100644 --- a/imperative/python/test/unit/data/test_pre_dataloader.py +++ b/imperative/python/test/unit/data/test_pre_dataloader.py @@ -209,7 +209,6 @@ def test_dataloader_parallel_timeout(): reason="dataloader do not support parallel on windows", ) def test_dataloader_parallel_worker_exception(): - print("in target") dataset = init_dataset() class FakeErrorTransform(Transform): @@ -217,7 +216,7 @@ def test_dataloader_parallel_worker_exception(): pass def apply(self, input): - y = x + 1 + raise RuntimeError("test raise error") return input dataloader = DataLoader( diff --git a/imperative/python/test/unit/data/test_transform.py b/imperative/python/test/unit/data/test_transform.py index 7fce8bf02..df9baade8 100644 --- a/imperative/python/test/unit/data/test_transform.py +++ b/imperative/python/test/unit/data/test_transform.py @@ -103,6 +103,7 @@ def test_Compose(): ) aug_data = t.apply_batch(generate_data()) aug_data_shape = [(a.shape, b.shape) for a, b in aug_data] - print(aug_data_shape) target_shape = [((3, 90, 70), label_shape)] * 4 - assert aug_data_shape == target_shape + assert aug_data_shape == target_shape, "aug {}, target {}".format( + aug_data_shape, target_shape + ) diff --git a/imperative/python/test/unit/functional/test_elemwise.py b/imperative/python/test/unit/functional/test_elemwise.py index 7e1c2680c..e4cae31c0 100644 --- a/imperative/python/test/unit/functional/test_elemwise.py +++ b/imperative/python/test/unit/functional/test_elemwise.py @@ -236,8 +236,7 @@ def test_empty_tensor(is_trace): elif nargs == 2: binary_func.append([op_name, op]) else: - print(nargs) - raise NotImplementedError + raise NotImplementedError("nargs {}".format(nargs)) def run_test(func, args, ref_shape, is_trace, sym=False): args = [tensor(t, dtype="float32") for t in args] @@ -248,8 +247,7 @@ def test_empty_tensor(is_trace): assert out.numpy().shape == ref_shape else: out = func(*args) - assert out.numpy().shape == ref_shape - print(out.numpy().shape) + assert out.numpy().shape == ref_shape, out.numpy().shape inps = [ np.array([]).astype("float32"), diff --git a/imperative/python/test/unit/functional/test_functional.py b/imperative/python/test/unit/functional/test_functional.py index bd73dbb81..4771562ec 100644 --- a/imperative/python/test/unit/functional/test_functional.py +++ b/imperative/python/test/unit/functional/test_functional.py @@ -922,8 +922,8 @@ def test_layer_norm(): def test_batchnorm2d_autocast(): """check amp's result is equal to manually converted result""" amp.enabled = True - tshape = (1, 224, 224, 3) - pshape = (1, 1, 1, 3) + tshape = (1, 3, 224, 224) + pshape = (1, 3, 1, 1) inp = tensor(np.random.randn(*tshape), dtype=np.float32) weight = tensor(np.ones(pshape, dtype=np.float32)) bias = tensor(np.zeros(pshape, dtype=np.float32)) @@ -948,7 +948,6 @@ def test_conv3d(): inp = tensor(np.ones((2, 2, 4, 4, 4), dtype=np.float32)) weight = tensor(np.ones((3, 2, 2, 2, 2), dtype=np.float32)) out = F.conv3d(inp, weight, None, 2, 0, 1, 1) - print(out.numpy().shape) np.testing.assert_equal( out.numpy(), np.ones((2, 3, 2, 2, 2), dtype=np.float32) * 16 ) diff --git a/src/opr/impl/dnn/batch_norm.cpp b/src/opr/impl/dnn/batch_norm.cpp index 72e70ec45..ce568ed78 100644 --- a/src/opr/impl/dnn/batch_norm.cpp +++ b/src/opr/impl/dnn/batch_norm.cpp @@ -230,15 +230,21 @@ void BatchNormForward::get_output_var_shape( for (size_t i = 0; i < 4; ++ i) { out_shape[i] = inp_shape[1]; } - out_shape[4] = {megdnn_opr()->get_reserve_in_bytes({inp_shape[0], input(0)->dtype()})}; if (!need_stats()) { out_shape[0] = out_shape[1] = {0}; } + if (inp_shape[0].is_empty()) { + out_shape[4] = {0}; + } else { + out_shape[4] = {megdnn_opr()->get_reserve_in_bytes({inp_shape[0], input(0)->dtype()})}; + } } size_t BatchNormForward::get_workspace_size_bytes( const TensorShapeArray &input_shapes, const TensorShapeArray &output_shapes) const { + if (input_shapes[0].is_empty()) + return 0; #define in(x) {input_shapes[x], input(x)->dtype()} #define out(x) {output_shapes[x], output(x)->dtype()} return megdnn_opr()->get_workspace_in_bytes( -- GitLab