Finish the test cases on ResNet-152

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__":