From 7c9046a6d26e37ecdbb2ed02059760b32a48670f Mon Sep 17 00:00:00 2001 From: Chaitanya Prakash Bapat Date: Wed, 23 Sep 2020 22:33:47 -0700 Subject: [PATCH] [v1.x] Nightly Large Tensor test cherrypicks (#19194) (#19215) * fixing batch_norm and layer_norm for large tensors (#17805) Co-authored-by: Rohit Kumar Srivastava * Fix nightly large_vector test caused by incorrect with_seed path (#18178) * add back the missing environment function Co-authored-by: Rohit Kumar Srivastava Co-authored-by: Rohit Kumar Srivastava Co-authored-by: Rohit Kumar Srivastava Co-authored-by: Rohit Kumar Srivastava --- src/operator/nn/batch_norm.cc | 2 +- src/operator/nn/layer_norm.cc | 2 +- tests/nightly/test_large_array.py | 3 +-- tests/nightly/test_large_vector.py | 3 +-- tests/python/unittest/common.py | 21 --------------------- 5 files changed, 4 insertions(+), 27 deletions(-) diff --git a/src/operator/nn/batch_norm.cc b/src/operator/nn/batch_norm.cc index 60f955399877..1bbdfa63160f 100644 --- a/src/operator/nn/batch_norm.cc +++ b/src/operator/nn/batch_norm.cc @@ -374,7 +374,7 @@ static bool BatchNormShape(const nnvm::NodeAttrs& attrs, : param.axis); CHECK_LT(channelAxis, dshape.ndim()) << "Channel axis out of range: " << param.axis; - const int channelCount = dshape[channelAxis]; + const index_t channelCount = dshape[channelAxis]; in_shape->at(batchnorm::kGamma) = mxnet::TShape(Shape1(channelCount)); in_shape->at(batchnorm::kBeta) = mxnet::TShape(Shape1(channelCount)); diff --git a/src/operator/nn/layer_norm.cc b/src/operator/nn/layer_norm.cc index c3ccd0d7a6bc..11178b358c2d 100644 --- a/src/operator/nn/layer_norm.cc +++ b/src/operator/nn/layer_norm.cc @@ -51,7 +51,7 @@ static bool LayerNormShape(const nnvm::NodeAttrs& attrs, CHECK(axis >= 0 && axis < dshape.ndim()) << "Channel axis out of range: axis=" << param.axis; - const int channelCount = dshape[axis]; + const index_t channelCount = dshape[axis]; SHAPE_ASSIGN_CHECK(*in_shape, layernorm::kGamma, diff --git a/tests/nightly/test_large_array.py b/tests/nightly/test_large_array.py index e26d73c327cc..d93d1835a1ec 100644 --- a/tests/nightly/test_large_array.py +++ b/tests/nightly/test_large_array.py @@ -27,9 +27,8 @@ from mxnet.test_utils import rand_ndarray, assert_almost_equal, rand_coord_2d, default_context, check_symbolic_forward, create_2d_tensor, get_identity_mat, get_identity_mat_batch from mxnet import gluon, nd -from common import with_seed, with_post_test_cleanup, assertRaises +from common import with_seed, assertRaises from mxnet.base import MXNetError -from nose.tools import with_setup import unittest # dimension constants diff --git a/tests/nightly/test_large_vector.py b/tests/nightly/test_large_vector.py index d4365391cf4e..4c81ddd7af79 100644 --- a/tests/nightly/test_large_vector.py +++ b/tests/nightly/test_large_vector.py @@ -27,9 +27,8 @@ from mxnet.test_utils import rand_ndarray, assert_almost_equal, rand_coord_2d, create_vector from mxnet import gluon, nd -from tests.python.unittest.common import with_seed, assertRaises +from common import with_seed, assertRaises from mxnet.base import MXNetError -from nose.tools import with_setup import unittest # dimension constants diff --git a/tests/python/unittest/common.py b/tests/python/unittest/common.py index cbddf0a327e6..9e51a3d46029 100644 --- a/tests/python/unittest/common.py +++ b/tests/python/unittest/common.py @@ -314,26 +314,6 @@ def teardown(): """ mx.nd.waitall() - -def with_post_test_cleanup(): - """ - Helper function that cleans up memory by releasing it from memory pool - Required especially by large tensor tests that have memory footprints in GBs. - """ - def test_helper(orig_test): - @make_decorator(orig_test) - def test_new(*args, **kwargs): - logger = default_logger() - try: - orig_test(*args, **kwargs) - except: - logger.info(test_msg) - raise - finally: - mx.nd.waitall() - mx.cpu().empty_cache() - - def with_environment(*args_): """ Helper function that takes a dictionary of environment variables and their @@ -349,7 +329,6 @@ def test_new(*args, **kwargs): return test_new return test_helper - def run_in_spawned_process(func, env, *args): """ Helper function to run a test in its own process.