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 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 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(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 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 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': {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)
