Fix cuDNN tests (keras-team#8189)

taehoonlee · fchollet · commit 7b35041e470e · 2017-10-20T10:47:23.000-07:00
Move `times = []` from the outer to the inner loop
Delete `clear_session()` which is redundant with `keras_test`
Reduce the example size under the 3x speed-up is satisfied
diff --git a/tests/keras/layers/cudnn_recurrent_test.py b/tests/keras/layers/cudnn_recurrent_test.py
@@ -122,46 +122,45 @@ def test_cudnn_rnn_canonical_to_params_gru():
 
 
 @keras_test
+@pytest.mark.parametrize('rnn_type', ['lstm', 'gru'], ids=['LSTM', 'GRU'])
 @pytest.mark.skipif((keras.backend.backend() != 'tensorflow'),
                     reason='Requires TensorFlow backend')
 @pytest.mark.skipif(not keras.backend.tensorflow_backend._get_available_gpus(),
                     reason='Requires GPU')
-def test_cudnn_rnn_timing():
+def test_cudnn_rnn_timing(rnn_type):
     input_size = 1000
     timesteps = 60
     units = 256
     num_samples = 10000
 
     times = []
-    for rnn_type in ['lstm', 'gru']:
-        for use_cudnn in [True, False]:
-            start_time = time.time()
-            inputs = keras.layers.Input(shape=(None, input_size))
-            if use_cudnn:
-                if rnn_type == 'lstm':
-                    layer = keras.layers.CuDNNLSTM(units)
-                else:
-                    layer = keras.layers.CuDNNGRU(units)
+    for use_cudnn in [True, False]:
+        start_time = time.time()
+        inputs = keras.layers.Input(shape=(None, input_size))
+        if use_cudnn:
+            if rnn_type == 'lstm':
+                layer = keras.layers.CuDNNLSTM(units)
             else:
-                if rnn_type == 'lstm':
-                    layer = keras.layers.LSTM(units)
-                else:
-                    layer = keras.layers.GRU(units)
-            outputs = layer(inputs)
+                layer = keras.layers.CuDNNGRU(units)
+        else:
+            if rnn_type == 'lstm':
+                layer = keras.layers.LSTM(units)
+            else:
+                layer = keras.layers.GRU(units)
+        outputs = layer(inputs)
 
-            model = keras.models.Model(inputs, outputs)
-            model.compile('sgd', 'mse')
+        model = keras.models.Model(inputs, outputs)
+        model.compile('sgd', 'mse')
 
-            x = np.random.random((num_samples, timesteps, input_size))
-            y = np.random.random((num_samples, units))
-            model.fit(x, y, epochs=4, batch_size=32)
+        x = np.random.random((num_samples, timesteps, input_size))
+        y = np.random.random((num_samples, units))
+        model.fit(x, y, epochs=4, batch_size=32)
 
-            times.append(time.time() - start_time)
+        times.append(time.time() - start_time)
 
-        speedup = times[1] / times[0]
-        print(rnn_type, 'speedup', speedup)
-        assert speedup > 3
-        keras.backend.clear_session()
+    speedup = times[1] / times[0]
+    print(rnn_type, 'speedup', speedup)
+    assert speedup > 3
 
 
 @keras_test
