from tensorflow.core.protobuf import config_pb2
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import server_lib
+from tensorflow.python.estimator import util
from tensorflow.python.util import compat_internal
from tensorflow.python.util.tf_export import tf_export
_USE_DEFAULT = object()
+_VALID_DEVICE_FN_ARGS = set(['op'])
# A list of the property names in RunConfig that the user is allowed to change.
_DEFAULT_REPLACEABLE_LIST = [
'keep_checkpoint_max',
'keep_checkpoint_every_n_hours',
'log_step_count_steps',
- 'train_distribute'
+ 'train_distribute',
+ 'device_fn'
]
_SAVE_CKPT_ERR = (
_validate('tf_random_seed', lambda seed: isinstance(seed, six.integer_types),
message='tf_random_seed must be integer.')
+ _validate('device_fn', lambda device_fn: six.callable(device_fn) and
+ set(util.fn_args(device_fn)) == _VALID_DEVICE_FN_ARGS,
+ message='device_fn must be callable with exactly'
+ ' one argument "op".')
+
class TaskType(object):
MASTER = 'master'
keep_checkpoint_max=5,
keep_checkpoint_every_n_hours=10000,
log_step_count_steps=100,
- train_distribute=None):
+ train_distribute=None,
+ device_fn=None):
"""Constructs a RunConfig.
All distributed training related properties `cluster_spec`, `is_chief`,
`tf.contrib.distribute.DistributionStrategy`. If specified,
then Estimator will distribute the user's model during training,
according to the policy specified by that strategy.
+ device_fn: A callable invoked for every `Operation` that takes the
+ `Operation` and returns the device string. If `None`, defaults to
+ the device function returned by `tf.train.replica_device_setter`
+ with round-robin strategy.
Raises:
ValueError: If both `save_checkpoints_steps` and `save_checkpoints_secs`
keep_checkpoint_max=keep_checkpoint_max,
keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours,
log_step_count_steps=log_step_count_steps,
- train_distribute=train_distribute)
+ train_distribute=train_distribute,
+ device_fn=device_fn)
self._init_distributed_setting_from_environment_var(tf_config)
return self._cluster_spec
@property
+ def device_fn(self):
+ """Returns the device_fn.
+
+ If device_fn is not `None`, it overrides the default
+ device function used in `Estimator`.
+ Otherwise the default one is used.
+ """
+ return self._device_fn
+
+ @property
def evaluation_master(self):
return self._evaluation_master
- `keep_checkpoint_max`,
- `keep_checkpoint_every_n_hours`,
- `log_step_count_steps`,
- - `train_distribute`.
+ - `train_distribute`,
+ - `device_fn`.
In addition, either `save_checkpoints_steps` or `save_checkpoints_secs`
can be set (should not be both).
_KEEP_CKPT_MAX_ERR = 'keep_checkpoint_max should be >= 0'
_KEEP_CKPT_HOURS_ERR = 'keep_checkpoint_every_n_hours should be > 0'
_TF_RANDOM_SEED_ERR = 'tf_random_seed must be integer'
+_DEVICE_FN_ERR = 'device_fn must be callable with exactly one argument "op".'
_ONE_CHIEF_ERR = 'The "cluster" in TF_CONFIG must have only one "chief" node.'
_ONE_MASTER_ERR = 'The "cluster" in TF_CONFIG must have only one "master" node.'
_INVALID_TASK_TYPE_FOR_EVAL_MASTER = (
self.assertEqual(5, config.keep_checkpoint_max)
self.assertEqual(10000, config.keep_checkpoint_every_n_hours)
self.assertIsNone(config.service)
+ self.assertIsNone(config.device_fn)
def test_model_dir(self):
empty_config = run_config_lib.RunConfig()
def test_replace_with_allowed_properties(self):
session_config = config_pb2.ConfigProto(allow_soft_placement=True)
+ device_fn = lambda op: "/cpu:0"
config = run_config_lib.RunConfig().replace(
tf_random_seed=11,
save_checkpoints_secs=14,
session_config=session_config,
keep_checkpoint_max=16,
- keep_checkpoint_every_n_hours=17)
+ keep_checkpoint_every_n_hours=17,
+ device_fn=device_fn)
self.assertEqual(11, config.tf_random_seed)
self.assertEqual(12, config.save_summary_steps)
self.assertEqual(14, config.save_checkpoints_secs)
self.assertEqual(session_config, config.session_config)
self.assertEqual(16, config.keep_checkpoint_max)
self.assertEqual(17, config.keep_checkpoint_every_n_hours)
+ self.assertEqual(device_fn, config.device_fn)
def test_replace_none_value(self):
config = run_config_lib.RunConfig().replace(
save_checkpoints_steps=None,
session_config=None,
keep_checkpoint_max=None,
- keep_checkpoint_every_n_hours=None)
+ keep_checkpoint_every_n_hours=None,
+ device_fn=None)
self.assertIsNone(config.tf_random_seed)
self.assertIsNone(config.model_dir)
self.assertIsNone(config.save_summary_steps)
self.assertIsNone(config.session_config)
self.assertIsNone(config.keep_checkpoint_max)
self.assertIsNone(config.keep_checkpoint_every_n_hours)
+ self.assertIsNone(config.device_fn)
def test_replace_with_disallowallowed_properties(self):
config = run_config_lib.RunConfig()
config.replace(keep_checkpoint_every_n_hours=0)
with self.assertRaisesRegexp(ValueError, _TF_RANDOM_SEED_ERR):
config.replace(tf_random_seed=1.0)
+ with self.assertRaisesRegexp(ValueError, _DEVICE_FN_ERR):
+ config.replace(device_fn=lambda x, y: 0)
def test_init_with_allowed_properties(self):
session_config = config_pb2.ConfigProto(allow_soft_placement=True)
+ device_fn = lambda op: "/cpu:0"
config = run_config_lib.RunConfig(
tf_random_seed=11,
save_checkpoints_secs=14,
session_config=session_config,
keep_checkpoint_max=16,
- keep_checkpoint_every_n_hours=17)
+ keep_checkpoint_every_n_hours=17,
+ device_fn=device_fn)
self.assertEqual(11, config.tf_random_seed)
self.assertEqual(12, config.save_summary_steps)
self.assertEqual(14, config.save_checkpoints_secs)
self.assertEqual(session_config, config.session_config)
self.assertEqual(16, config.keep_checkpoint_max)
self.assertEqual(17, config.keep_checkpoint_every_n_hours)
+ self.assertEqual(device_fn, config.device_fn)
def test_init_none_value(self):
config = run_config_lib.RunConfig(
save_checkpoints_steps=None,
session_config=None,
keep_checkpoint_max=None,
- keep_checkpoint_every_n_hours=None)
+ keep_checkpoint_every_n_hours=None,
+ device_fn=None)
self.assertIsNone(config.tf_random_seed)
self.assertIsNone(config.model_dir)
self.assertIsNone(config.save_summary_steps)
self.assertIsNone(config.session_config)
self.assertIsNone(config.keep_checkpoint_max)
self.assertIsNone(config.keep_checkpoint_every_n_hours)
+ self.assertIsNone(config.device_fn)
def test_init_invalid_values(self):
with self.assertRaisesRegexp(ValueError, _MODEL_DIR_ERR):
run_config_lib.RunConfig(keep_checkpoint_every_n_hours=0)
with self.assertRaisesRegexp(ValueError, _TF_RANDOM_SEED_ERR):
run_config_lib.RunConfig(tf_random_seed=1.0)
+ with self.assertRaisesRegexp(ValueError, _DEVICE_FN_ERR):
+ run_config_lib.RunConfig(device_fn=lambda x: "/cpu:0")
class RunConfigDistributedSettingTest(test.TestCase):
mtype: "<type \'property\'>"
}
member {
+ name: "device_fn"
+ mtype: "<type \'property\'>"
+ }
+ member {
name: "evaluation_master"
mtype: "<type \'property\'>"
}
}
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'model_dir\', \'tf_random_seed\', \'save_summary_steps\', \'save_checkpoints_steps\', \'save_checkpoints_secs\', \'session_config\', \'keep_checkpoint_max\', \'keep_checkpoint_every_n_hours\', \'log_step_count_steps\', \'train_distribute\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'100\', \'<object object instance>\', \'<object object instance>\', \'None\', \'5\', \'10000\', \'100\', \'None\'], "
+ argspec: "args=[\'self\', \'model_dir\', \'tf_random_seed\', \'save_summary_steps\', \'save_checkpoints_steps\', \'save_checkpoints_secs\', \'session_config\', \'keep_checkpoint_max\', \'keep_checkpoint_every_n_hours\', \'log_step_count_steps\', \'train_distribute\', \'device_fn\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'100\', \'<object object instance>\', \'<object object instance>\', \'None\', \'5\', \'10000\', \'100\', \'None\', \'None\'], "
}
member_method {
name: "replace"