def _wrapped_model(*args):
"""A concrete tf.function that wraps the model's call function."""
# When given a single input, Keras models will call the model on the tensor
# rather than a list consisting of the single tensor.
inputs = args[0] if len(input_signature) == 1 else list(args)
outputs_list = nest.flatten(model(inputs=inputs))
try:
output_names = model.output_names
except AttributeError:
from tensorflow.python.keras.engine import training_utils # pylint: disable=g-import-not-at-top
output_names = training_utils.generic_output_names(outputs_list)
return {name: output for name, output in zip(output_names, outputs_list)}
def compile(self,
optimizer,
loss,
n_replicas,
metrics=None,
target_tensors=None):
if any(arg is None for arg in (optimizer, loss, n_replicas)):
raise ValueError('The arg is None')
# validate losses
# ...
# validate optimizer
# ...
self.n_replicas = n_replicas
metrics = metrics or []
self._stateful_metrics_names = _stateful_metrics_names(metrics)
# set config for `tf.keras.optimizers.Optimizer` subclasses
if isinstance(optimizer, tf.keras.optimizers.Optimizer):
config = optimizer.get_config()
optimizer_config = {'class_name': config['name'], 'config': config}
elif isinstance(optimizer, str):
# TODO: Add test for string optimizer. Check if config == {}
# then `_set_hyper(placeholder)` works as expected.
optimizer_config = {'class_name': optimizer, 'config': {}}
elif (isinstance(optimizer, abc.Mapping) and 'class_name' in optimizer
and 'config' in optimizer):
optimizer_config = optimizer
else:
raise NotImplementedError("Not recognized optimizer.")
optimizer_config['config'].pop('name', None)
train_attrs = {i: {} for i in range(self.n_replicas)}
for i in range(self.n_replicas):
with tf.device(training_utils.gpu_device_name(i)):
# build model and the state of hyperparameters
with tf.variable_scope('model_%d' % i):
hp = training_utils.HyperParamState()
train_attrs[i]['hp_state'] = hp
# Each model has its own input placeholder, meaning that the
# feed values are fed `n_replica` times.
# TODO: Implement this by using the same input for each one
# of the model. This could reduce overhead of copying to
# GPUs. In particular, the behaviour of
# `prepare_input_output_tensors()` and values should be modified.
model = self._model_builder_fn(hp)
with tf.variable_scope('loss_%d' % i):
outputs = model.outputs
output_names = keras_train_utils.generic_output_names(
outputs)
loss_functions = keras_train_utils.prepare_loss_functions(
loss, output_names)
# Set placeholders instead of actual values for each possible
# hyperparameterd of the optimizer. All float values in `_hyper`
# could be exchaned between different replicas. By default,
# if the value is not exchanged the default value is fed. No need
# to take care of feeding values that are not being exchanged.
opt = tf.keras.optimizers.get(optimizer_config)
for n, v in opt._hyper.items():
if isinstance(v, float):
opt._set_hyper(n, hp.get_hparam(n, default_value=v))
# Each replica will have now its own metric class.
# For example, if `tf.keras.metrics.Precision()`
# is passed to metrics we will duplicate this class
# `n_replica` times.
# TODO: right now if the initialized class is passed it is not
# used. A new class is created instead. This could be improved
# by using this class for the first replica. Same could be done
# of initialized optimizer.
compiled_metrics = []
for m in metrics:
if isinstance(m, str):
m = m
elif isinstance(m, tf.keras.metrics.Metric):
args_kwargs = training_utils._infer_init_args_kwargs(m)
_ = args_kwargs.pop('name', None)
m = m.__class__(**args_kwargs)
elif callable(m):
m = m
else:
raise ValueError('unexpected metric', m)
compiled_metrics.append(m)
train_attrs[i].update({
'model': model,
'loss_functions': loss_functions,
'optimizer': opt,
'compiled_metrics': compiled_metrics
})
self._train_attrs = train_attrs
self.inputs = list(
itertools.chain(
*[train_attrs[i]['model'].inputs for i in range(n_replicas)]))
self.outputs = list(
itertools.chain(
*[train_attrs[i]['model'].outputs for i in range(n_replicas)]))
self._is_compiled = True