def wrap_optimizer(self, optimizer: Any) -> Any:
"""
This should be used to wrap optimizer objects immediately after they have
been created. Users should use the output of this wrapper as the new instance
of their optimizer. For example, if users create their optimizer within
``build_estimator()``, they should call ``optimizer = wrap_optimizer(optimzer)``
prior to passing the optimizer into their Estimator.
"""
self.optimizer_initialized = True
if not self.hvd_config.use:
return optimizer
check.check_false(
isinstance(optimizer, str),
"Please specify an optimizer object instead of using a string name.",
)
hvd.require_horovod_type(
"tensorflow", "EstimatorContext.wrap_optimizer was called.")
use_compression = self.hvd_config.fp16_compression
optimizer = hvd.DistributedOptimizer(
optimizer,
compression=hvd.compression.Compression.fp16
if use_compression else hvd.compression.Compression.none,
aggregation_frequency=self.hvd_config.aggregation_frequency,
average_aggregated_gradients=self.hvd_config.
average_aggregated_gradients,
)
logging.debug(
"Initialized optimizer for distributed and optimized parallel training."
)
return optimizer
def cache_dataset(
self, dataset_id: str, dataset_version: str, shuffle: bool, skip_shuffle_at_epoch_end: bool,
) -> Callable:
# Perform lazy initialization of storage so that if users are not
# using data layer, we are not creating unused directories.
self._configure_storage()
if self._training:
# We only check the training cacheable for re-use, because for EstimatorTrial
# it's possible that the validation cacheable is called every time validation
# is performed.
check.check_false(
self._decorator_used,
"Pleas use both `@context.experimental.cache_train_dataset(dataset_name, "
"dataset_version)` and `@context.experimental.cache_validation_dataset("
"dataset_name, dataset_version)` exactly once.",
)
self._decorator_used = True
dataset_version += "_train" if self._training else "_val"
def _wrap(make_dataset_fn: Callable) -> Callable:
@functools.wraps(make_dataset_fn)
def _decorated_fn(*args: Any, **kwargs: Any) -> Any:
@self._storage.cacheable( # type: ignore
dataset_id=dataset_id, dataset_version=dataset_version,
)
def make_dataset() -> yogadl.DataRef:
return make_dataset_fn(*args, **kwargs)
logging.info(f"Preparing dataset: {dataset_id}:{dataset_version}.")
logging.debug(
f"Calling make dataset for: {dataset_id}:{dataset_version} "
f"with following start_offset: {self._offset}, "
f"shuffle: {shuffle} shuffle_seed: {self._shuffle_seed} "
f"shard_rank: {self._shard_rank}, world size: {self._num_shards} "
f"training: {self._training}."
)
stream_from_cache = make_dataset().stream(
start_offset=self._offset,
shuffle=shuffle,
skip_shuffle_at_epoch_end=skip_shuffle_at_epoch_end,
shuffle_seed=self._shuffle_seed,
shard_rank=self._shard_rank,
num_shards=self._num_shards,
drop_shard_remainder=True if self._training else False,
)
self._dataset_length = len(stream_from_cache)
logging.info(f"Dataset {dataset_id}:{dataset_version} preparation finished.")
return tensorflow.make_tf_dataset(stream_from_cache)
return _decorated_fn
return _wrap
def patch(args: argparse.Namespace) -> None:
check_false(args.all and args.agent_id)
if not (args.all or args.agent_id):
print("Error: must specify exactly one of `--all` or agent_id")
sys.exit(1)
if args.agent_id:
agent_ids = [args.agent_id]
else:
r = api.get(args.master, "agents")
agent_ids = sorted(local_id(a) for a in r.json().keys())
for agent_id in agent_ids:
path = "agents/{}/slots".format(agent_id)
headers = {"Content-Type": "application/merge-patch+json"}
payload = {"enabled": enabled}
api.patch(args.master, path, body=payload, headers=headers)
status = "Disabled" if not enabled else "Enabled"
print("{} agent {}".format(status, agent_id))
def wrap_optimizer(self, optimizer: Any) -> Any:
"""
This should be used to wrap optimizer objects immediately after they have
been created. Users should use the output of this wrapper as the new instance
of their optimizer. For example, if users create their optimizer within
``build_estimator()``, they should call ``optimizer = wrap_optimizer(optimzer)``
prior to passing the optimizer into their Estimator.
"""
if not self.env.managed_training:
return optimizer
self.optimizer_initialized = True
if not self.hvd_config.use:
return optimizer
check.check_false(
isinstance(optimizer, str),
"Please specify an optimizer object instead of using a string name.",
)
hvd.require_horovod_type("tensorflow", "EstimatorContext.wrap_optimizer was called.")
use_compression = self.hvd_config.fp16_compression
# The signature of our horovod optimizer changed after we rebased onto 0.21.
hvd_sig = inspect.signature(hvd.DistributedOptimizer)
horovod_kwargs = {
"compression": hvd.compression.Compression.fp16
if use_compression
else hvd.compression.Compression.none,
"average_aggregated_gradients": self.hvd_config.average_aggregated_gradients,
}
if "aggregation_frequency" in hvd_sig.parameters:
horovod_kwargs["aggregation_frequency"] = self.hvd_config.aggregation_frequency
else:
horovod_kwargs["backward_passes_per_step"] = self.hvd_config.aggregation_frequency
optimizer = hvd.DistributedOptimizer(optimizer, **horovod_kwargs)
logging.debug("Initialized optimizer for distributed and optimized parallel training.")
return optimizer