From 78af5c7e2cf9e7a09acfceb368d6bd3818405353 Mon Sep 17 00:00:00 2001 From: Igor Saprykin Date: Mon, 5 Feb 2018 17:24:31 -0800 Subject: [PATCH] Correctly treat "devices=/gpu:0" argument of replicate_model_fn. At the moment if "devices=/GPU:0" are specified by the user, then variables are going to be placed on the GPU. However, if "devices=/gpu:0" are given, then they are going to be placed on the CPU. Instead, the latter case should be equivalent to the former case. PiperOrigin-RevId: 184612823 --- .../python/estimator/replicate_model_fn.py | 2 +- .../python/estimator/replicate_model_fn_test.py | 28 ++++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/tensorflow/contrib/estimator/python/estimator/replicate_model_fn.py b/tensorflow/contrib/estimator/python/estimator/replicate_model_fn.py index c9153c9..dfae034 100644 --- a/tensorflow/contrib/estimator/python/estimator/replicate_model_fn.py +++ b/tensorflow/contrib/estimator/python/estimator/replicate_model_fn.py @@ -195,7 +195,7 @@ def _replicate_model_fn_with_mode( if not devices: devices = _get_local_devices('GPU') or _get_local_devices('CPU') - is_a_single_gpu_case = len(devices) == 1 and 'GPU' in devices[0] + is_a_single_gpu_case = len(devices) == 1 and 'GPU' in devices[0].upper() consolidation_device = devices[0] if is_a_single_gpu_case else '/CPU:0' ps_devices = [consolidation_device] diff --git a/tensorflow/contrib/estimator/python/estimator/replicate_model_fn_test.py b/tensorflow/contrib/estimator/python/estimator/replicate_model_fn_test.py index 6936f8a..ab117e6 100644 --- a/tensorflow/contrib/estimator/python/estimator/replicate_model_fn_test.py +++ b/tensorflow/contrib/estimator/python/estimator/replicate_model_fn_test.py @@ -451,6 +451,34 @@ class ReplicateModelTest(test_util.TensorFlowTestCase): _ = replicate_model_fn.replicate_model_fn(self.model_fn, losses.Reduction.NONE) + def test_places_on_gpu_with_upper_case_spelling(self): + features = np.array([[0.01], [0.002]]) + labels = np.array([[0.01], [0.02]]) + + with self.test_session(): + replicated_model_fn = replicate_model_fn.replicate_model_fn( + self.model_fn, devices=['/GPU:0']) + _ = replicated_model_fn( + features, labels, model_fn_lib.ModeKeys.TRAIN, self.params) + + with variable_scope.variable_scope('', reuse=True): + c = variable_scope.get_variable('c', dtype=dtypes.float64) + self.assertEqual('/device:GPU:0', c.device) + + def test_places_on_gpu_with_lower_case_spelling(self): + features = np.array([[0.01], [0.002]]) + labels = np.array([[0.01], [0.02]]) + + with self.test_session(): + replicated_model_fn = replicate_model_fn.replicate_model_fn( + self.model_fn, devices=['/gpu:0']) + _ = replicated_model_fn( + features, labels, model_fn_lib.ModeKeys.TRAIN, self.params) + + with variable_scope.variable_scope('', reuse=True): + c = variable_scope.get_variable('c', dtype=dtypes.float64) + self.assertEqual('/device:GPU:0', c.device) + class ReplicateAcrossASingleDeviceWithoutTowerOptimizer( test_util.TensorFlowTestCase): -- 2.7.4