def testMutableHashTableIsLocal(self): with tf.device(estimator._get_replica_device_setter(tf.contrib.learn.RunConfig(num_ps_replicas=0))): default_val = tf.constant([-1, -1], tf.int64) table = tf.contrib.lookup.MutableHashTable(tf.string, tf.int64, default_val) input_string = tf.constant(["brain", "salad", "tank"]) output = table.lookup(input_string) self.assertDeviceEqual("", table._table_ref.device) self.assertDeviceEqual("", output.device)
def testVariablesAreLocal(self): with tf.device(estimator._get_replica_device_setter(tf.contrib.learn.RunConfig(num_ps_replicas=0))): v = tf.Variable([1, 2]) w = tf.Variable([2, 1]) a = v + w self.assertDeviceEqual("", v.device) self.assertDeviceEqual("", v.initializer.device) self.assertDeviceEqual("", w.device) self.assertDeviceEqual("", w.initializer.device) self.assertDeviceEqual("", a.device)
def testVariablesAreLocal(self): with tf.device(estimator._get_replica_device_setter(num_ps_replicas=0)): v = tf.Variable([1, 2]) w = tf.Variable([2, 1]) a = v + w self.assertDeviceEqual('', v.device) self.assertDeviceEqual('', v.initializer.device) self.assertDeviceEqual('', w.device) self.assertDeviceEqual('', w.initializer.device) self.assertDeviceEqual('', a.device)
def testMutableHashTableIsLocal(self): with tf.device(estimator._get_replica_device_setter(num_ps_replicas=0)): default_val = tf.constant([-1, -1], tf.int64) table = tf.contrib.lookup.MutableHashTable(tf.string, tf.int64, default_val) input_string = tf.constant(['brain', 'salad', 'tank']) output = table.lookup(input_string) self.assertDeviceEqual('', table._table_ref.device) self.assertDeviceEqual('', output.device)
def testMutableHashTableIsLocal(self): with tf.device( estimator._get_replica_device_setter(num_ps_replicas=0)): default_val = tf.constant([-1, -1], tf.int64) table = tf.contrib.lookup.MutableHashTable(tf.string, tf.int64, default_val) input_string = tf.constant(['brain', 'salad', 'tank']) output = table.lookup(input_string) self.assertDeviceEqual('', table._table_ref.device) self.assertDeviceEqual('', output.device)
def testVariablesAreOnPs(self): with tf.device(estimator._get_replica_device_setter(num_ps_replicas=1)): v = tf.Variable([1, 2]) w = tf.Variable([2, 1]) a = v + w self.assertDeviceEqual('/job:ps/task:0', v.device) self.assertDeviceEqual('/job:ps/task:0', v.initializer.device) self.assertDeviceEqual('/job:ps/task:0', w.device) self.assertDeviceEqual('/job:ps/task:0', w.initializer.device) self.assertDeviceEqual('/job:worker', a.device)
def testMutableHashTableIsLocal(self): with ops.device( estimator._get_replica_device_setter(run_config.RunConfig())): default_val = constant_op.constant([-1, -1], dtypes.int64) table = lookup.MutableHashTable(dtypes.string, dtypes.int64, default_val) input_string = constant_op.constant(['brain', 'salad', 'tank']) output = table.lookup(input_string) self.assertDeviceEqual('', table._table_ref.device) self.assertDeviceEqual('', output.device)
def testVariablesAreLocal(self): with ops.device( estimator._get_replica_device_setter(run_config.RunConfig())): v = variables_lib.Variable([1, 2]) w = variables_lib.Variable([2, 1]) a = v + w self.assertDeviceEqual('', v.device) self.assertDeviceEqual('', v.initializer.device) self.assertDeviceEqual('', w.device) self.assertDeviceEqual('', w.initializer.device) self.assertDeviceEqual('', a.device)
def testVariablesAreLocal(self): with tf.device(estimator._get_replica_device_setter( tf.contrib.learn.RunConfig())): v = tf.Variable([1, 2]) w = tf.Variable([2, 1]) a = v + w self.assertDeviceEqual('', v.device) self.assertDeviceEqual('', v.initializer.device) self.assertDeviceEqual('', w.device) self.assertDeviceEqual('', w.initializer.device) self.assertDeviceEqual('', a.device)
def testVariablesAreLocal(self): with tf.device( estimator._get_replica_device_setter(num_ps_replicas=0)): v = tf.Variable([1, 2]) w = tf.Variable([2, 1]) a = v + w self.assertDeviceEqual('', v.device) self.assertDeviceEqual('', v.initializer.device) self.assertDeviceEqual('', w.device) self.assertDeviceEqual('', w.initializer.device) self.assertDeviceEqual('', a.device)
def testVariablesAreOnPs(self): with tf.device(estimator._get_replica_device_setter( tf.contrib.learn.RunConfig(num_ps_replicas=1))): v = tf.Variable([1, 2]) w = tf.Variable([2, 1]) a = v + w self.assertDeviceEqual('/job:ps/task:0', v.device) self.assertDeviceEqual('/job:ps/task:0', v.initializer.device) self.assertDeviceEqual('/job:ps/task:0', w.device) self.assertDeviceEqual('/job:ps/task:0', w.initializer.device) self.assertDeviceEqual('/job:worker', a.device)
def testMutableHashTableIsOnPs(self): tf_config = {"cluster": {tf.contrib.learn.TaskType.PS: ["fake_ps_0"]}} with tf.test.mock.patch.dict("os.environ", {"TF_CONFIG": json.dumps(tf_config)}): config = tf.contrib.learn.RunConfig() with tf.device(estimator._get_replica_device_setter(config)): default_val = tf.constant([-1, -1], tf.int64) table = tf.contrib.lookup.MutableHashTable(tf.string, tf.int64, default_val) input_string = tf.constant(["brain", "salad", "tank"]) output = table.lookup(input_string) self.assertDeviceEqual("/job:ps/task:0", table._table_ref.device) self.assertDeviceEqual("/job:ps/task:0", output.device)
def testTaskIsSetOnWorkerWhenJobNameIsSet(self): with tf.device( estimator._get_replica_device_setter( tf.contrib.learn.RunConfig( num_ps_replicas=1, job_name='worker', task=3))): v = tf.Variable([1, 2]) w = tf.Variable([2, 1]) a = v + w self.assertDeviceEqual('/job:ps/task:0', v.device) self.assertDeviceEqual('/job:ps/task:0', v.initializer.device) self.assertDeviceEqual('/job:ps/task:0', w.device) self.assertDeviceEqual('/job:ps/task:0', w.initializer.device) self.assertDeviceEqual('/job:worker/task:3', a.device)
def testMutableHashTableIsOnPs(self): tf_config = {'cluster': {run_config.TaskType.PS: ['fake_ps_0']}} with test.mock.patch.dict('os.environ', {'TF_CONFIG': json.dumps(tf_config)}): config = run_config.RunConfig() with ops.device(estimator._get_replica_device_setter(config)): default_val = constant_op.constant([-1, -1], dtypes.int64) table = lookup.MutableHashTable(dtypes.string, dtypes.int64, default_val) input_string = constant_op.constant(['brain', 'salad', 'tank']) output = table.lookup(input_string) self.assertDeviceEqual('/job:ps/task:0', table._table_ref.device) self.assertDeviceEqual('/job:ps/task:0', output.device)
def testVariablesAreOnPs(self): tf_config = {"cluster": {tf.contrib.learn.TaskType.PS: ["fake_ps_0"]}} with tf.test.mock.patch.dict("os.environ", {"TF_CONFIG": json.dumps(tf_config)}): config = tf.contrib.learn.RunConfig() with tf.device(estimator._get_replica_device_setter(config)): v = tf.Variable([1, 2]) w = tf.Variable([2, 1]) a = v + w self.assertDeviceEqual("/job:ps/task:0", v.device) self.assertDeviceEqual("/job:ps/task:0", v.initializer.device) self.assertDeviceEqual("/job:ps/task:0", w.device) self.assertDeviceEqual("/job:ps/task:0", w.initializer.device) self.assertDeviceEqual("/job:worker", a.device)
def testVariablesAreOnPs(self): tf_config = {'cluster': {run_config.TaskType.PS: ['fake_ps_0']}} with test.mock.patch.dict('os.environ', {'TF_CONFIG': json.dumps(tf_config)}): config = run_config.RunConfig() with ops.device(estimator._get_replica_device_setter(config)): v = variables_lib.Variable([1, 2]) w = variables_lib.Variable([2, 1]) a = v + w self.assertDeviceEqual('/job:ps/task:0', v.device) self.assertDeviceEqual('/job:ps/task:0', v.initializer.device) self.assertDeviceEqual('/job:ps/task:0', w.device) self.assertDeviceEqual('/job:ps/task:0', w.initializer.device) self.assertDeviceEqual('/job:worker', a.device)
def testVariablesAreOnPs(self): tf_config = {'cluster': {tf.contrib.learn.TaskType.PS: ['fake_ps_0']}} with tf.test.mock.patch.dict('os.environ', {'TF_CONFIG': json.dumps(tf_config)}): config = tf.contrib.learn.RunConfig() with tf.device(estimator._get_replica_device_setter(config)): v = tf.Variable([1, 2]) w = tf.Variable([2, 1]) a = v + w self.assertDeviceEqual('/job:ps/task:0', v.device) self.assertDeviceEqual('/job:ps/task:0', v.initializer.device) self.assertDeviceEqual('/job:ps/task:0', w.device) self.assertDeviceEqual('/job:ps/task:0', w.initializer.device) self.assertDeviceEqual('/job:worker', a.device)
def testTaskIsSetOnWorkerWhenJobNameIsSet(self): tf_config = { "cluster": {tf.contrib.learn.TaskType.PS: ["fake_ps_0"]}, "task": {"type": tf.contrib.learn.TaskType.WORKER, "index": 3}, } with tf.test.mock.patch.dict("os.environ", {"TF_CONFIG": json.dumps(tf_config)}): config = tf.contrib.learn.RunConfig() with tf.device(estimator._get_replica_device_setter(config)): v = tf.Variable([1, 2]) w = tf.Variable([2, 1]) a = v + w self.assertDeviceEqual("/job:ps/task:0", v.device) self.assertDeviceEqual("/job:ps/task:0", v.initializer.device) self.assertDeviceEqual("/job:ps/task:0", w.device) self.assertDeviceEqual("/job:ps/task:0", w.initializer.device) self.assertDeviceEqual("/job:worker/task:3", a.device)
def __init__(self, model_fn=None, model_dir=None, config=None, params=None): # Create a run configuration. if config is None: self._config = RunConfig() logging.info("Using default config.") else: if not isinstance(config, RunConfig): raise ValueError("config must be an instance of RunConfig, " "received {}.".format(config)) self._config = config logging.info("Using config: {}".format(vars(self._config))) if (model_dir is not None) and (self._config.model_dir is not None): if model_dir != self._config.model_dir: # pylint: disable=g-doc-exception raise ValueError( "model_dir are set both in constructor and RunConfig, but with " "different values. In constructor: '{}', in RunConfig: " "'{}' ".format(model_dir, self._config.model_dir)) self._model_dir = model_dir or self._config.model_dir or generate_model_dir( ) if self._config.model_dir is None: self._config = self._config.replace(model_dir=self._model_dir) if self._config.session_config is None: self._session_config = config_pb2.ConfigProto( allow_soft_placement=True) else: self._session_config = self._config.session_config # Set device function depending if there are replicas or not. self._device_fn = _get_replica_device_setter(self._config) self._graph = None self._verify_model_fn_args(model_fn, params) self._model_fn = model_fn self._params = params or {}
def __init__(self, model_fn, model_dir=None, config=None, params=None): # Create a run configuration. if config is None: self._config = RunConfig() logging.info("Using default config.") else: if not isinstance(config, RunConfig): raise ValueError("config must be an instance of RunConfig, " "received {}.".format(config)) self._config = config if(model_dir is not None) and (self._config.model_dir is not None): if model_dir != self._config.model_dir: # pylint: disable=g-doc-exception raise ValueError( "model_dir are set both in constructor and RunConfig, but with " "different values. In constructor: '{}', in RunConfig: " "'{}' ".format(model_dir, self._config.model_dir)) self._model_dir = model_dir or self._config.model_dir or generate_model_dir() if self._config.model_dir is None: self._config = self._config.replace(model_dir=self._model_dir) logging.info("Using config: {}".format(vars(self._config))) if self._config.session_config is None: self._session_config = config_pb2.ConfigProto(allow_soft_placement=True) else: self._session_config = self._config.session_config # Set device function depending if there are replicas or not. self._device_fn = _get_replica_device_setter(self._config) self._graph = None self._verify_model_fn_args(model_fn, params) self._model_fn = model_fn self._params = params or {}
def testTaskIsSetOnWorkerWhenJobNameIsSet(self): tf_config = { 'cluster': { tf.contrib.learn.TaskType.PS: ['fake_ps_0'] }, 'task': { 'type': tf.contrib.learn.TaskType.WORKER, 'index': 3 } } with tf.test.mock.patch.dict('os.environ', {'TF_CONFIG': json.dumps(tf_config)}): config = tf.contrib.learn.RunConfig() with tf.device(estimator._get_replica_device_setter(config)): v = tf.Variable([1, 2]) w = tf.Variable([2, 1]) a = v + w self.assertDeviceEqual('/job:ps/task:0', v.device) self.assertDeviceEqual('/job:ps/task:0', v.initializer.device) self.assertDeviceEqual('/job:ps/task:0', w.device) self.assertDeviceEqual('/job:ps/task:0', w.initializer.device) self.assertDeviceEqual('/job:worker/task:3', a.device)
def testTaskIsSetOnWorkerWhenJobNameIsSet(self): tf_config = { 'cluster': { run_config.TaskType.PS: ['fake_ps_0'] }, 'task': { 'type': run_config.TaskType.WORKER, 'index': 3 } } with test.mock.patch.dict('os.environ', {'TF_CONFIG': json.dumps(tf_config)}): config = run_config.RunConfig() with ops.device(estimator._get_replica_device_setter(config)): v = variables_lib.Variable([1, 2]) w = variables_lib.Variable([2, 1]) a = v + w self.assertDeviceEqual('/job:ps/task:0', v.device) self.assertDeviceEqual('/job:ps/task:0', v.initializer.device) self.assertDeviceEqual('/job:ps/task:0', w.device) self.assertDeviceEqual('/job:ps/task:0', w.initializer.device) self.assertDeviceEqual('/job:worker/task:3', a.device)