diff --git a/tests/python/unittest/test_memory_opt.py b/tests/python/unittest/test_memory_opt.py index a8f735c827c1..65dab0516e03 100644 --- a/tests/python/unittest/test_memory_opt.py +++ b/tests/python/unittest/test_memory_opt.py @@ -141,6 +141,27 @@ def test_fc(): "Not all operator nodes have been verified on the mirror stage" +def grep_exec_memory_consumption(exec): + # Grep the memory consumption (in MB) from the executor debug string. + # + # It is important to note that, due to various reasons, the memory + # consumption reported by the executor debug string might be very different + # when compared with the real numbers reported by nvidia-smi. These reasons + # include: + # - Allocations by the CUDA Library (e.g., cuDNN, cuBLAS) + # - Fragmentation (of the MXNet Memory Allocator and cudaMalloc) + exec_debug_str = exec.debug_str().split('\n') + + import re # We will be using regular expressions for grepping the model + # memory consumption. + alloc_line_pattern = re.compile("Total \d+ MB allocated") + for line in exec_debug_str(): + if alloc_line_pattern.match(line) is not None: + return int(line.split()[1]) + assert False, "Unable to gerp the memory consumption numbers from the executor " \ + "debug string: %s" % exec_debug_str + + def test_resnet152(): # Verify the memory allocation behavior on ResNet-152, the state-of-the-art # model used for image classification. @@ -151,16 +172,20 @@ def test_resnet152(): return from importlib import import_module resnet_mod = import_module('symbols_resnet') - resnet_50 = resnet_mod.get_symbol(num_classes=1000, - num_layers=152, - image_shape='3,224,224') + resnet_152 = resnet_mod.get_symbol(num_classes=1000, + num_layers=152, + image_shape='3,224,224') # We do the binding twice, one with the memory optimizations and one without. # It is expected that the memory consumption of the former should be roughly # half of that of the latter. - memory_opt_exec = resnet_50.simple_bind(mx.cpu(), 'write', data=(16, 3, 224, 224)) + memory_opt_exec = resnet_152.simple_bind(mx.cpu(), 'write', + data=(32, 3, 224, 224)) os.environ["MXNET_MEMORY_OPT"] = '0' - no_opt_exec = resnet_50.simple_bind(mx.cpu(), 'write', data=(16, 3, 224, 224)) - print(memory_opt_exec.debug_str(), no_opt_exec.debug_str()) + no_opt_exec = resnet_152.simple_bind(mx.cpu(), 'write', data=(32, 3, 224, 224)) + assert grep_exec_memory_consumption(memory_opt_exec) / \ + grep_exec_memory_consumption(no_opt_exec) < 0.6, \ + "The ratio between the memory consumption with the memory optimizations " \ + "enabled and disabled is expected to be smaller than 0.6" if __name__ == "__main__":