Fix tests to deal with OOM error on CNTK (#6)

Author: taehoonlee

tests/applications_test.py

@@ -28,6 +28,8 @@
 from keras.preprocessing import image
 from keras import backend
 
+from multiprocessing import Process, Queue
+
 
 MOBILENET_LIST = [(mobilenet.MobileNet, 1024),
                   (mobilenet_v2.MobileNetV2, 1280)]
@@ -38,43 +40,76 @@
                (nasnet.NASNetLarge, 4032)]
 
 
-@keras_test
-def _test_application_basic(app, last_dim=1000, module=None):
-    if module is not None:
-        weights = 'imagenet'
-    else:
-        weights = None
-    model = app(weights=weights)
-    assert model.output_shape == (None, last_dim)
-    if module is None:
-        return
-
-    img_path = 'tests/data/elephant.jpg'
-    target_size = tuple(model.input_shape[1: 3])
+def _get_elephant(target_size):
     # For models that don't include a Flatten step,
     # the default is to accept variable-size inputs
     # even when loading ImageNet weights (since it is possible).
     # In this case, default to 299x299.
     if target_size[0] is None:
         target_size = (299, 299)
-    img = image.load_img(img_path, target_size=target_size)
+    img = image.load_img('tests/data/elephant.jpg',
+                         target_size=tuple(target_size))
     x = image.img_to_array(img)
-    x = np.expand_dims(x, axis=0)
+    return np.expand_dims(x, axis=0)
+
+
+def _get_output_shape(model_fn, preprocess_input=None):
+    if backend.backend() == 'cntk':
+        # Create model in a subprocess so that
+        # the memory consumed by InceptionResNetV2 will be
+        # released back to the system after this test
+        # (to deal with OOM error on CNTK backend).
+        # TODO: remove the use of multiprocessing from these tests
+        # once a memory clearing mechanism
+        # is implemented in the CNTK backend.
+        def target(queue):
+            model = model_fn()
+            if preprocess_input is None:
+                queue.put(model.output_shape)
+            else:
+                x = _get_elephant(model.input_shape[1:3])
+                x = preprocess_input(x)
+                queue.put((model.output_shape, model.predict(x)))
+        queue = Queue()
+        p = Process(target=target, args=(queue,))
+        p.start()
+        p.join()
+        # The error in a subprocess won't propagate
+        # to the main process, so we check if the model
+        # is successfully created by checking if the output shape
+        # has been put into the queue
+        assert not queue.empty(), 'Model creation failed.'
+        return queue.get_nowait()
+    else:
+        model = model_fn()
+        if preprocess_input is None:
+            return model.output_shape
+        else:
+            x = _get_elephant(model.input_shape[1:3])
+            x = preprocess_input(x)
+            return (model.output_shape, model.predict(x))
+
 
-    preprocess_input = getattr(module, 'preprocess_input')
-    decode_predictions = getattr(module, 'decode_predictions')
-    x = preprocess_input(x)
+@keras_test
+def _test_application_basic(app, last_dim=1000, module=None):
+    if module is None:
+        output_shape = _get_output_shape(lambda: app(weights=None))
+        assert output_shape == (None, None, None, last_dim)
+    else:
+        output_shape, preds = _get_output_shape(
+            lambda: app(weights='imagenet'), module.preprocess_input)
+        assert output_shape == (None, last_dim)
 
-    preds = model.predict(x)
-    names = [p[1] for p in decode_predictions(preds)[0]]
-    # Test correct label is in top 3 (weak correctness test).
-    assert 'African_elephant' in names[:3]
+        names = [p[1] for p in module.decode_predictions(preds)[0]]
+        # Test correct label is in top 3 (weak correctness test).
+        assert 'African_elephant' in names[:3]
 
 
 @keras_test
 def _test_application_notop(app, last_dim):
-    model = app(weights=None, include_top=False)
-    assert model.output_shape == (None, None, None, last_dim)
+    output_shape = _get_output_shape(
+        lambda: app(weights=None, include_top=False))
+    assert output_shape == (None, None, None, last_dim)
 
 
 @keras_test
@@ -83,23 +118,26 @@ def _test_application_variable_input_channels(app, last_dim):
         input_shape = (1, None, None)
     else:
         input_shape = (None, None, 1)
-    model = app(weights=None, include_top=False, input_shape=input_shape)
-    assert model.output_shape == (None, None, None, last_dim)
+    output_shape = _get_output_shape(
+        lambda: app(weights=None, include_top=False, input_shape=input_shape))
+    assert output_shape == (None, None, None, last_dim)
 
     if backend.image_data_format() == 'channels_first':
         input_shape = (4, None, None)
     else:
         input_shape = (None, None, 4)
-    model = app(weights=None, include_top=False, input_shape=input_shape)
-    assert model.output_shape == (None, None, None, last_dim)
+    output_shape = _get_output_shape(
+        lambda: app(weights=None, include_top=False, input_shape=input_shape))
+    assert output_shape == (None, None, None, last_dim)
 
 
 @keras_test
 def _test_app_pooling(app, last_dim):
-    model = app(weights=None,
-                include_top=False,
-                pooling=random.choice(['avg', 'max']))
-    assert model.output_shape == (None, last_dim)
+    output_shape = _get_output_shape(
+        lambda: app(weights=None,
+                    include_top=False,
+                    pooling=random.choice(['avg', 'max'])))
+    assert output_shape == (None, last_dim)
 
 
 def test_resnet50():