also set intra_op_parallelism_threads (keras-team#12254)

* also set intra_op_parallelism_threads * add test * tensorflow only * skip if not tensorflow
yumios · Feb 27, 2019 · 2944988 · 2944988
1 parent 050e8cb
commit 2944988
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/keras/backend/tensorflow_backend.py b/keras/backend/tensorflow_backend.py
@@ -186,6 +186,7 @@ def get_session():
             else:
                 num_thread = int(os.environ.get('OMP_NUM_THREADS'))
                 config = tf.ConfigProto(intra_op_parallelism_threads=num_thread,
+                                        inter_op_parallelism_threads=num_thread,
                                         allow_soft_placement=True)
             _SESSION = tf.Session(config=config)
         session = _SESSION

diff --git a/tests/keras/backend/backend_test.py b/tests/keras/backend/backend_test.py
@@ -2118,6 +2118,15 @@ def test_clip_supports_tensor_arguments(self):
         assert np.allclose(K.eval(K.clip(x, min_value, max_value)),
                            np.asarray([-5., -4., 0., 4., 9.], dtype=np.float32))
 
+    @pytest.mark.skipif(K.backend() != 'tensorflow',
+                        reason='This test is for tensorflow parallelism.')
+    def test_tensorflow_session_parallelism_settings(self, monkeypatch):
+        for threads in [0, 1, 4]:
+            K.clear_session()
+            monkeypatch.setenv('OMP_NUM_THREADS', str(threads))
+            cfg = K.get_session()._config
+            assert cfg.intra_op_parallelism_threads == threads
+            assert cfg.inter_op_parallelism_threads == threads
 
 if __name__ == '__main__':
     pytest.main([__file__])