def create_ensemble_metrics(metric_fn,
use_tpu=False,
features=None,
labels=None,
estimator_spec=None,
architecture=None):
"""Creates an instance of the _EnsembleMetrics class.
Args:
metric_fn: A function which should obey the following signature:
- Args: can only have following three arguments in any order:
* predictions: Predictions `Tensor` or dict of `Tensor` created by given
`Head`.
* features: Input `dict` of `Tensor` objects created by `input_fn` which
is given to `estimator.evaluate` as an argument.
* labels: Labels `Tensor` or dict of `Tensor` (for multi-head) created
by `input_fn` which is given to `estimator.evaluate` as an argument.
- Returns: Dict of metric results keyed by name. Final metrics are a union
of this and `estimator`s existing metrics. If there is a name conflict
between this and `estimator`s existing metrics, this will override the
existing one. The values of the dict are the results of calling a metric
function, namely a `(metric_tensor, update_op)` tuple.
use_tpu: Whether to use TPU-specific variable sharing logic.
features: Input `dict` of `Tensor` objects.
labels: Labels `Tensor` or a dictionary of string label name to `Tensor`
(for multi-head).
estimator_spec: The `EstimatorSpec` created by a `Head` instance.
architecture: `_Architecture` object.
Returns:
An instance of _EnsembleMetrics.
"""
if not estimator_spec:
estimator_spec = tf_compat.v1.estimator.tpu.TPUEstimatorSpec(
mode=tf.estimator.ModeKeys.EVAL,
loss=tf.constant(2.),
predictions=None,
eval_metrics=None)
if not use_tpu:
estimator_spec = estimator_spec.as_estimator_spec()
if not architecture:
architecture = _Architecture(None, None)
metrics = _EnsembleMetrics(use_tpu=use_tpu)
metrics.create_eval_metrics(features, labels, estimator_spec, metric_fn,
architecture)
return metrics
def test_ensemble_metrics(self):
architecture = _Architecture("test_ensemble_candidate")
architecture.add_subnetwork(iteration_number=0, builder_name="b_0_0")
architecture.add_subnetwork(iteration_number=0, builder_name="b_0_1")
architecture.add_subnetwork(iteration_number=1, builder_name="b_1_0")
architecture.add_subnetwork(iteration_number=2, builder_name="b_2_0")
metrics = _EnsembleMetrics()
metrics.create_eval_metrics(self._features, self._labels,
self._estimator_spec, self._metric_fn,
architecture)
with self.test_session() as sess:
actual = _run_metrics(sess, metrics.eval_metrics_tuple())
serialized_arch_proto = actual["architecture/adanet/ensembles"]
expected_arch_string = b"| b_0_0 | b_0_1 | b_1_0 | b_2_0 |"
self.assertIn(expected_arch_string, serialized_arch_proto)
def build_ensemble_spec(self,
name,
candidate,
ensembler,
subnetwork_specs,
summary,
features,
mode,
iteration_step,
iteration_number,
labels=None,
previous_ensemble_spec=None):
"""Builds an `_EnsembleSpec` with the given `adanet.ensemble.Candidate`.
Args:
name: The string name of the ensemble. Typically the name of the builder
that returned the given `Subnetwork`.
candidate: The `adanet.ensemble.Candidate` for this spec.
ensembler: The :class:`adanet.ensemble.Ensembler` to use to ensemble a
group of subnetworks.
subnetwork_specs: Iterable of `_SubnetworkSpecs` for this iteration.
summary: A `_ScopedSummary` instance for recording ensemble summaries.
features: Input `dict` of `Tensor` objects.
mode: Estimator `ModeKeys` indicating training, evaluation, or inference.
iteration_step: Integer `Tensor` representing the step since the beginning
of the current iteration, as opposed to the global step.
iteration_number: Integer current iteration number.
labels: Labels `Tensor` or a dictionary of string label name to `Tensor`
(for multi-head).
previous_ensemble_spec: Link the rest of the `_EnsembleSpec` from
iteration t-1. Used for creating the subnetwork train_op.
Returns:
An `_EnsembleSpec` instance.
"""
with tf.variable_scope("ensemble_{}".format(name)):
architecture = _Architecture(candidate.name)
previous_subnetworks = []
subnetwork_builders = []
previous_ensemble = None
if previous_ensemble_spec:
previous_ensemble = previous_ensemble_spec.ensemble
previous_architecture = previous_ensemble_spec.architecture
keep_indices = range(len(previous_ensemble.subnetworks))
if len(candidate.subnetwork_builders) == 1 and previous_ensemble:
# Prune previous ensemble according to the subnetwork.Builder for
# backwards compatibility.
tf.logging.warn(
"Using an `adanet.subnetwork.Builder#prune_previous_ensemble` "
"is deprecated. Please use a custom `adanet.ensemble.Strategy` "
"instead.")
subnetwork_builder = candidate.subnetwork_builders[0]
keep_indices = subnetwork_builder.prune_previous_ensemble(
previous_ensemble)
for i, builder in enumerate(previous_ensemble_spec.subnetwork_builders):
if i not in keep_indices:
continue
if builder not in candidate.previous_ensemble_subnetwork_builders:
continue
previous_subnetworks.append(previous_ensemble.subnetworks[i])
subnetwork_builders.append(builder)
architecture.add_subnetwork(*previous_architecture.subnetworks[i])
for builder in candidate.subnetwork_builders:
architecture.add_subnetwork(iteration_number, builder.name)
subnetwork_builders.append(builder)
subnetwork_map = {s.builder.name: s.subnetwork for s in subnetwork_specs}
subnetworks = [
subnetwork_map[s.name] for s in candidate.subnetwork_builders
]
ensemble_scope = tf.get_variable_scope()
before_var_list = tf.trainable_variables()
with summary.current_scope(), _monkey_patch_context(
iteration_step_scope=ensemble_scope,
scoped_summary=summary,
trainable_vars=[]):
ensemble = ensembler.build_ensemble(
subnetworks,
previous_ensemble_subnetworks=previous_subnetworks,
features=features,
labels=labels,
logits_dimension=self._head.logits_dimension,
training=mode == tf.estimator.ModeKeys.TRAIN,
iteration_step=iteration_step,
summary=summary,
previous_ensemble=previous_ensemble)
ensemble_var_list = _new_trainable_variables(before_var_list)
estimator_spec = _create_estimator_spec(
self._head, features, labels, mode, ensemble.logits, self._use_tpu)
ensemble_loss = estimator_spec.loss
adanet_loss = None
if mode != tf.estimator.ModeKeys.PREDICT:
# TODO: Support any kind of Ensemble. Use a moving average of
# their train loss for the 'adanet_loss'.
if not isinstance(ensemble, ComplexityRegularized):
raise ValueError(
"Only ComplexityRegularized ensembles are supported.")
adanet_loss = estimator_spec.loss + ensemble.complexity_regularization
ensemble_metrics = _EnsembleMetrics()
if mode == tf.estimator.ModeKeys.EVAL:
ensemble_metrics.create_eval_metrics(
features=features,
labels=labels,
estimator_spec=estimator_spec,
metric_fn=self._metric_fn,
architecture=architecture)
if mode == tf.estimator.ModeKeys.TRAIN:
with summary.current_scope():
summary.scalar("loss", estimator_spec.loss)
# Create train ops for training subnetworks and ensembles.
train_op = None
if mode == tf.estimator.ModeKeys.TRAIN:
# Note that these mixture weights are on top of the last_layer of the
# subnetwork constructed in TRAIN mode, which means that dropout is
# still applied when the mixture weights are being trained.
ensemble_scope = tf.get_variable_scope()
with tf.variable_scope("train_mixture_weights"):
with summary.current_scope(), _monkey_patch_context(
iteration_step_scope=ensemble_scope,
scoped_summary=summary,
trainable_vars=ensemble_var_list):
# For backwards compatibility.
subnetwork_builder = candidate.subnetwork_builders[0]
old_train_op_fn = getattr(subnetwork_builder,
"build_mixture_weights_train_op", None)
if callable(old_train_op_fn):
tf.logging.warn(
"The `build_mixture_weights_train_op` method is deprecated. "
"Please use the `Ensembler#build_train_op` instead.")
train_op = _to_train_op_spec(
subnetwork_builder.build_mixture_weights_train_op(
loss=adanet_loss,
var_list=ensemble_var_list,
logits=ensemble.logits,
labels=labels,
iteration_step=iteration_step,
summary=summary))
else:
train_op = _to_train_op_spec(
ensembler.build_train_op(
ensemble=ensemble,
loss=adanet_loss,
var_list=ensemble_var_list,
labels=labels,
iteration_step=iteration_step,
summary=summary,
previous_ensemble=previous_ensemble))
return _EnsembleSpec(
name=name,
architecture=architecture,
subnetwork_builders=subnetwork_builders,
ensemble=ensemble,
predictions=estimator_spec.predictions,
loss=ensemble_loss,
adanet_loss=adanet_loss,
train_op=train_op,
eval_metrics=ensemble_metrics.eval_metrics_tuple(),
export_outputs=estimator_spec.export_outputs)
def build_ensemble_spec(self, name, candidate, ensembler, subnetwork_specs,
summary, features, mode, iteration_number, labels,
my_ensemble_index, previous_ensemble_spec,
previous_iteration_checkpoint):
"""Builds an `_EnsembleSpec` with the given `adanet.ensemble.Candidate`.
Args:
name: The string name of the ensemble. Typically the name of the builder
that returned the given `Subnetwork`.
candidate: The `adanet.ensemble.Candidate` for this spec.
ensembler: The :class:`adanet.ensemble.Ensembler` to use to ensemble a
group of subnetworks.
subnetwork_specs: Iterable of `_SubnetworkSpecs` for this iteration.
summary: A `_ScopedSummary` instance for recording ensemble summaries.
features: Input `dict` of `Tensor` objects.
mode: Estimator `ModeKeys` indicating training, evaluation, or inference.
iteration_number: Integer current iteration number.
labels: Labels `Tensor` or a dictionary of string label name to `Tensor`
(for multi-head).
my_ensemble_index: An integer holding the index of the ensemble in the
candidates list of AdaNet.
previous_ensemble_spec: Link the rest of the `_EnsembleSpec` from
iteration t-1. Used for creating the subnetwork train_op.
previous_iteration_checkpoint: `tf.train.Checkpoint` for iteration t-1.
Returns:
An `_EnsembleSpec` instance.
"""
with tf_compat.v1.variable_scope("ensemble_{}".format(name)):
step = tf_compat.v1.get_variable(
"step",
shape=[],
initializer=tf_compat.v1.zeros_initializer(),
trainable=False,
dtype=tf.int64)
# Convert to tensor so that users cannot mutate it.
step_tensor = tf.convert_to_tensor(value=step)
with summary.current_scope():
summary.scalar("iteration_step/adanet/iteration_step",
step_tensor)
replay_indices = []
if previous_ensemble_spec:
replay_indices = copy.copy(
previous_ensemble_spec.architecture.replay_indices)
if my_ensemble_index is not None:
replay_indices.append(my_ensemble_index)
architecture = _Architecture(candidate.name,
ensembler.name,
replay_indices=replay_indices)
previous_subnetworks = []
previous_subnetwork_specs = []
subnetwork_builders = []
previous_ensemble = None
if previous_ensemble_spec:
previous_ensemble = previous_ensemble_spec.ensemble
previous_architecture = previous_ensemble_spec.architecture
keep_indices = range(len(previous_ensemble.subnetworks))
if len(candidate.subnetwork_builders
) == 1 and previous_ensemble:
# Prune previous ensemble according to the subnetwork.Builder for
# backwards compatibility.
subnetwork_builder = candidate.subnetwork_builders[0]
prune_previous_ensemble = getattr(
subnetwork_builder, "prune_previous_ensemble", None)
if callable(prune_previous_ensemble):
logging.warn(
"Using an `adanet.subnetwork.Builder#prune_previous_ensemble` "
"is deprecated. Please use a custom `adanet.ensemble.Strategy` "
"instead.")
keep_indices = prune_previous_ensemble(
previous_ensemble)
for i, builder in enumerate(
previous_ensemble_spec.subnetwork_builders):
if i not in keep_indices:
continue
if builder not in candidate.previous_ensemble_subnetwork_builders:
continue
previous_subnetworks.append(
previous_ensemble.subnetworks[i])
previous_subnetwork_specs.append(
previous_ensemble_spec.subnetwork_specs[i])
subnetwork_builders.append(builder)
architecture.add_subnetwork(
*previous_architecture.subnetworks[i])
for builder in candidate.subnetwork_builders:
architecture.add_subnetwork(iteration_number, builder.name)
subnetwork_builders.append(builder)
subnetwork_spec_map = {s.builder.name: s for s in subnetwork_specs}
relevant_subnetwork_specs = [
subnetwork_spec_map[s.name]
for s in candidate.subnetwork_builders
]
ensemble_scope = tf_compat.v1.get_variable_scope()
old_vars = _get_current_vars()
with summary.current_scope(), _monkey_patch_context(
iteration_step_scope=ensemble_scope,
scoped_summary=summary,
trainable_vars=[]):
ensemble = ensembler.build_ensemble(
subnetworks=[
s.subnetwork for s in relevant_subnetwork_specs
],
previous_ensemble_subnetworks=previous_subnetworks,
features=features,
labels=labels,
logits_dimension=self._head.logits_dimension,
training=mode == tf.estimator.ModeKeys.TRAIN,
iteration_step=step_tensor,
summary=summary,
previous_ensemble=previous_ensemble,
previous_iteration_checkpoint=previous_iteration_checkpoint
)
estimator_spec = _create_estimator_spec(self._head, features,
labels, mode,
ensemble.logits,
self._use_tpu)
ensemble_loss = estimator_spec.loss
adanet_loss = None
if mode != tf.estimator.ModeKeys.PREDICT:
adanet_loss = estimator_spec.loss
# Add ensembler specific loss
if isinstance(ensemble, ensemble_lib.ComplexityRegularized):
adanet_loss += ensemble.complexity_regularization
predictions = estimator_spec.predictions
export_outputs = estimator_spec.export_outputs
if (self._export_subnetwork_logits and export_outputs
and subnetwork_spec_map):
first_subnetwork_logits = list(
subnetwork_spec_map.values())[0].subnetwork.logits
if isinstance(first_subnetwork_logits, dict):
for head_name in first_subnetwork_logits.keys():
subnetwork_logits = {
subnetwork_name:
subnetwork_spec.subnetwork.logits[head_name]
for subnetwork_name, subnetwork_spec in
subnetwork_spec_map.items()
}
export_outputs.update({
"{}_{}".format(
_EnsembleBuilder._SUBNETWORK_LOGITS_EXPORT_SIGNATURE, head_name):
tf.estimator.export.PredictOutput(
subnetwork_logits)
})
else:
subnetwork_logits = {
subnetwork_name: subnetwork_spec.subnetwork.logits
for subnetwork_name, subnetwork_spec in
subnetwork_spec_map.items()
}
export_outputs.update({
_EnsembleBuilder._SUBNETWORK_LOGITS_EXPORT_SIGNATURE:
tf.estimator.export.PredictOutput(subnetwork_logits)
})
if (self._export_subnetwork_last_layer and export_outputs
and subnetwork_spec_map and list(
subnetwork_spec_map.values())[0].subnetwork.last_layer
is not None):
first_subnetwork_last_layer = list(
subnetwork_spec_map.values())[0].subnetwork.last_layer
if isinstance(first_subnetwork_last_layer, dict):
for head_name in first_subnetwork_last_layer.keys():
subnetwork_last_layer = {
subnetwork_name:
subnetwork_spec.subnetwork.last_layer[head_name]
for subnetwork_name, subnetwork_spec in
subnetwork_spec_map.items()
}
export_outputs.update({
"{}_{}".format(
_EnsembleBuilder._SUBNETWORK_LAST_LAYER_EXPORT_SIGNATURE, head_name):
tf.estimator.export.PredictOutput(
subnetwork_last_layer)
})
else:
subnetwork_last_layer = {
subnetwork_name: subnetwork_spec.subnetwork.last_layer
for subnetwork_name, subnetwork_spec in
subnetwork_spec_map.items()
}
export_outputs.update({
_EnsembleBuilder._SUBNETWORK_LAST_LAYER_EXPORT_SIGNATURE:
tf.estimator.export.PredictOutput(
subnetwork_last_layer)
})
if ensemble.predictions and predictions:
predictions.update(ensemble.predictions)
if ensemble.predictions and export_outputs:
export_outputs.update({
k: tf.estimator.export.PredictOutput(v)
for k, v in ensemble.predictions.items()
})
ensemble_metrics = _EnsembleMetrics(use_tpu=self._use_tpu)
if mode == tf.estimator.ModeKeys.EVAL:
ensemble_metrics.create_eval_metrics(
features=features,
labels=labels,
estimator_spec=estimator_spec,
metric_fn=self._metric_fn,
architecture=architecture)
if mode == tf.estimator.ModeKeys.TRAIN:
with summary.current_scope():
summary.scalar("loss", estimator_spec.loss)
ensemble_trainable_vars = _get_current_vars(
diffbase=old_vars)["trainable"]
# Create train ops for training subnetworks and ensembles.
train_op = None
if mode == tf.estimator.ModeKeys.TRAIN:
# Note that these mixture weights are on top of the last_layer of the
# subnetwork constructed in TRAIN mode, which means that dropout is
# still applied when the mixture weights are being trained.
ensemble_scope = tf_compat.v1.get_variable_scope()
with tf_compat.v1.variable_scope("train_mixture_weights"):
with summary.current_scope(), _monkey_patch_context(
iteration_step_scope=ensemble_scope,
scoped_summary=summary,
trainable_vars=ensemble_trainable_vars):
# For backwards compatibility.
subnetwork_builder = candidate.subnetwork_builders[0]
old_train_op_fn = getattr(
subnetwork_builder,
"build_mixture_weights_train_op", None)
if callable(old_train_op_fn):
logging.warn(
"The `build_mixture_weights_train_op` method is deprecated. "
"Please use the `Ensembler#build_train_op` instead."
)
train_op = _to_train_op_spec(
subnetwork_builder.
build_mixture_weights_train_op(
loss=adanet_loss,
var_list=ensemble_trainable_vars,
logits=ensemble.logits,
labels=labels,
iteration_step=step_tensor,
summary=summary))
else:
train_op = _to_train_op_spec(
ensembler.build_train_op(
ensemble=ensemble,
loss=adanet_loss,
var_list=ensemble_trainable_vars,
labels=labels,
iteration_step=step_tensor,
summary=summary,
previous_ensemble=previous_ensemble))
new_vars = _get_current_vars(diffbase=old_vars)
# Sort our dictionary by key to remove non-determinism of variable order.
new_vars = collections.OrderedDict(sorted(new_vars.items()))
# Combine all trainable, global and savable variables into a single list.
ensemble_variables = sum(new_vars.values(), []) + [step]
return _EnsembleSpec(name=name,
architecture=architecture,
subnetwork_builders=subnetwork_builders,
subnetwork_specs=previous_subnetwork_specs +
relevant_subnetwork_specs,
ensemble=ensemble,
predictions=predictions,
step=step,
variables=ensemble_variables,
loss=ensemble_loss,
adanet_loss=adanet_loss,
train_op=train_op,
eval_metrics=ensemble_metrics,
export_outputs=export_outputs)
