"""An AdaNet estimator implementation in Tensorflow using a single graph.
Copyright 2018 The AdaNet Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import contextlib
import errno
import inspect
import os
import time
from absl import logging
from adanet import distributed as distributed_lib
from adanet import ensemble as ensemble_lib
from adanet import tf_compat
from adanet.core.architecture import _Architecture
from adanet.core.candidate import _CandidateBuilder
from adanet.core.ensemble_builder import _EnsembleBuilder
from adanet.core.ensemble_builder import _SubnetworkManager
from adanet.core.iteration import _Iteration
from adanet.core.iteration import _IterationBuilder
from adanet.core.report_accessor import _ReportAccessor
from adanet.core.summary import _ScopedSummary
from adanet.core.summary import _ScopedSummaryV2
from adanet.core.summary import _TPUScopedSummary
from adanet.core.timer import _CountDownTimer
from adanet.distributed.devices import monkey_patch_default_variable_placement_strategy
import numpy as np
import six
import tensorflow.compat.v2 as tf
from typing import Any, Callable, Dict, Optional, Sequence, Text # (b/144172555) pylint:disable=unused-import
class _StopAfterTrainingHook(tf_compat.SessionRunHook):
"""Hook that requests stop once iteration is over."""
def __init__(self, iteration, after_fn):
# type: (_Iteration, Callable[[], None]) -> None
"""Initializes a `_StopAfterTrainingHook`.
Args:
iteration: An `_Iteration` instance.
after_fn: A function to call after training stopped.
Returns:
A `_StopAfterTrainingHook` instance.
"""
self._iteration = iteration
self._after_fn = after_fn
def before_run(self, run_context):
"""See `SessionRunHook`."""
self._stop_if_is_over(run_context)
def after_run(self, run_context, run_values):
"""See `SessionRunHook`."""
self._stop_if_is_over(run_context)
def _stop_if_is_over(self, run_context):
"""Signals the monitored session to step when the iteration is over."""
if not self._iteration.train_manager.is_over():
return
logging.info("Now stopping iteration %d training", self._iteration.number)
run_context.request_stop()
self._after_fn()
class _SummaryV2SaverHook(tf_compat.SessionRunHook):
"""A hook that writes summaries to the appropriate log directory on disk."""
def __init__(self, summaries, save_steps=None, save_secs=None):
"""Initializes a `SummaryV2SaverHook` for writing TF 2 summaries.
Args:
summaries: List of `_ScopedSummaryV2` instances.
save_steps: `int`, save summaries every N steps. Exactly one of
`save_secs` and `save_steps` should be set.
save_secs: `int`, save summaries every N seconds.
"""
self._summaries = summaries
self._summary_ops = []
self._writer_init_ops = []
self._timer = tf_compat.v1.train.SecondOrStepTimer(
every_secs=save_secs, every_steps=save_steps)
def begin(self):
self._next_step = None
self._global_step_tensor = tf_compat.v1.train.get_global_step()
for summary in self._summaries:
assert isinstance(summary, _ScopedSummaryV2)
writer = tf_compat.v2.summary.create_file_writer(summary.logdir)
with writer.as_default():
for summary_fn, tensor in summary.summary_tuples():
self._summary_ops.append(
summary_fn(tensor, step=tf.compat.v1.train.get_global_step()))
self._writer_init_ops.append(writer.init())
def after_create_session(self, session, coord):
session.run(self._writer_init_ops)
def before_run(self, run_context):
requests = {"global_step": self._global_step_tensor}
self._request_summary = (
self._next_step is None or
self._timer.should_trigger_for_step(self._next_step))
if self._request_summary:
requests["summary"] = self._summary_ops
return tf_compat.SessionRunArgs(requests)
def after_run(self, run_context, run_values):
stale_global_step = run_values.results["global_step"]
global_step = stale_global_step + 1
if self._next_step is None or self._request_summary:
global_step = run_context.session.run(self._global_step_tensor)
if self._request_summary:
self._timer.update_last_triggered_step(global_step)
self._next_step = global_step + 1
def end(self, session):
# TODO: Run writer.flush() at Session end.
# Currently disabled because the flush op crashes between iterations.
return
class _EvalMetricSaverHook(tf_compat.SessionRunHook):
"""A hook for writing candidate evaluation metrics as summaries to disk."""
def __init__(self, name, kind, eval_metrics, output_dir):
# type: (Text, Text, Any, Text) -> None
"""Initializes a `_EvalMetricSaverHook` instance.
Args:
name: String name of candidate owner of these metrics.
kind: The kind of candidate that the metrics belong to (e.g. subnetwork).
eval_metrics: Tuple of (metric_fn, tensors) which returns a dict of metric
results keyed by name. The values of the dict are the results of calling
a metric function, namely a `(metric_tensor, update_op)` tuple.
`metric_tensor` should be evaluated without any impact on state
(typically is a pure computation based on variables.). For example, it
should not trigger the `update_op` or require any input fetching.
output_dir: Directory for writing evaluation summaries.
Returns:
An `_EvalMetricSaverHook` instance.
"""
self._name = name
self._kind = kind
self._eval_metrics = eval_metrics
self._output_dir = output_dir
def begin(self):
"""See `SessionRunHook`."""
# The metric_fn is called with tf.placeholders to simply read the value of
# the metric variables. The metrics themselves are computed as a result of
# being returned in the EstimatorSpec by _adanet_model_fn.
metric_fn, tensors = self._eval_metrics.eval_metrics_tuple()
tensors = [tf_compat.v1.placeholder(t.dtype, t.shape) for t in tensors]
eval_metric_ops = metric_fn(*tensors)
self._eval_metric_tensors = {}
for key in sorted(eval_metric_ops):
value = tf_compat.metric_op(eval_metric_ops[key])
self._eval_metric_tensors[key] = value[0]
def _dict_to_str(self, dictionary):
"""Get a `str` representation of a `dict`.
Args:
dictionary: The `dict` to be represented as `str`.
Returns:
A `str` representing the `dictionary`.
"""
return ", ".join(
"{} = {}".format(k, v) for k, v in sorted(dictionary.items()))
def end(self, session):
"""See `SessionRunHook`."""
# Forked from tensorflow/python/estimator/estimator.py function called
# _write_dict_to_summary.
current_global_step = tf_compat.v1.train.get_global_step()
eval_dict, current_global_step = session.run(
(self._eval_metric_tensors, current_global_step))
logging.info("Saving %s '%s' dict for global step %d: %s", self._kind,
self._name, current_global_step, self._dict_to_str(eval_dict))
summary_writer = tf_compat.v1.summary.FileWriterCache.get(self._output_dir)
summary_proto = tf_compat.v1.summary.Summary()
for key in eval_dict:
value = eval_dict[key]
if isinstance(value, (np.float32, float)):
summary_proto.value.add(tag=key, simple_value=float(value))
elif isinstance(value, six.binary_type):
summ = tf_compat.v1.summary.Summary.FromString(value)
for i, _ in enumerate(summ.value):
summ.value[i].tag = "{}/{}".format(key, i)
summary_proto.value.extend(summ.value)
else:
logging.warn(
"Skipping summary for %s, must be a float, np.float32, "
"or a serialized string of Summary.", key)
summary_writer.add_summary(summary_proto, current_global_step)
summary_writer.flush()
# Note(b/137672676): Do not explicitly call summary_writer.close() here.
# This will cause eval summaries to not be written out after the first time
# in continuous evals.
class _OverwriteCheckpointHook(tf_compat.SessionRunHook):
"""Hook to overwrite the latest checkpoint with next iteration variables."""
def __init__(self, current_iteration, iteration_number_tensor,
previous_iteration_vars, config, enable_v2_checkpoint):
"""Initializes an _OverwriteCheckpointHook instance.
Args:
current_iteration: Current `_Iteration` object.
iteration_number_tensor: Int variable `Tensor` storing the current
iteration number.
previous_iteration_vars: Variables to restore from the previous iteration
before overwriting the checkpoint.
config: The Estimator's RunConfig object.
enable_v2_checkpoint: Whether `tf.train.Checkpoint` is used for
checkpointing.
"""
self._current_iteration = current_iteration
self._iteration_number = current_iteration.number
self._iteration_number_tensor = iteration_number_tensor
self._previous_iteration_vars = previous_iteration_vars
self._model_dir = config.model_dir
self._checkpoint_state = tf.train.get_checkpoint_state(self._model_dir)
self._keep_checkpoint_max = config.keep_checkpoint_max
self._enable_v2_checkpoint = enable_v2_checkpoint
self._update_op = None
self._overwrite_saver = None
self._checkpoint_overwritten = False
def begin(self):
"""Creates the savers and adds ops needed for overwriting the checkpoint.
Two savers are created, a restore saver which is passed the variables from
the previous iteration to restore, and an overwrite saver which will
actually overwrite the checkpoint.
"""
from tensorflow.python.training.tracking import graph_view # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
if self._enable_v2_checkpoint:
prev_checkpoint = self._current_iteration.previous_iteration.checkpoint
self._status = prev_checkpoint.restore(
self._checkpoint_state.model_checkpoint_path)
# Because we prune the previous iteration's candidates, only a subset of
# the variables present in the checkpoint will be used. Assert they are
# restored.
self._status.expect_partial().assert_existing_objects_matched()
self._overwrite_saver = tf_compat.v1.train.Saver(
var_list=graph_view.ObjectGraphView(
self._current_iteration.checkpoint).frozen_saveable_objects(),
sharded=True,
max_to_keep=self._keep_checkpoint_max)
else:
self._restore_saver = tf_compat.v1.train.Saver(
sharded=True, var_list=self._previous_iteration_vars)
# Note: self._iteration_number already contains the value of the next
# iteration since _OverwriteCheckpointHook should only execute during the
# graph growing phase.
self._update_op = self._iteration_number_tensor.assign(
self._iteration_number)
self._overwrite_saver = tf_compat.v1.train.Saver(
sharded=True, max_to_keep=self._keep_checkpoint_max)
self._overwrite_saver.recover_last_checkpoints(
self._checkpoint_state.all_model_checkpoint_paths)
def before_run(self, run_context):
"""Overwrites checkpoint before any calls to session.run().
This is to ensure that the values of the variables in the overwritten
checkpoint match those in the pevious iteration checkpoint.
Args:
run_context: The tf.train.SessionRunContext passed to the hook.
"""
if not self._checkpoint_overwritten:
session = run_context.session
if self._enable_v2_checkpoint:
self._status.initialize_or_restore(session)
else:
self._restore_saver.restore(
session, self._checkpoint_state.model_checkpoint_path)
session.run(self._update_op)
checkpoint_path = os.path.join(self._model_dir, "increment.ckpt")
logging.info(
"Overwriting checkpoint with new graph for iteration %d to %s-%d",
self._current_iteration.number, checkpoint_path,
self._current_iteration.number)
# Specify global_step=self._iteration_number to append the iteration
# number to the checkpoint name, e.g. <model_dir>/increment.ckpt-1.
self._overwrite_saver.save(
session, checkpoint_path, global_step=self._current_iteration.number)
self._checkpoint_overwritten = True
def _copy_recursively(source, destination):
"""Copies a directory and its content.
Args:
source: Source directory.
destination: Destination directory.
"""
for src_dir, _, src_files in tf.io.gfile.walk(source):
dst_dir = os.path.join(destination, os.path.relpath(src_dir, source))
if not tf.io.gfile.exists(dst_dir):
tf.io.gfile.makedirs(dst_dir)
for src_file in src_files:
tf.io.gfile.copy(
os.path.join(src_dir, src_file),
os.path.join(dst_dir, src_file),
overwrite=True)
class _GraphGrowingHookDecorator(tf_compat.SessionRunHook):
"""Decorates a SessionRunHook to only run begin() and end() methods."""
def __init__(self, hook):
# type: (tf_compat.SessionRunHook) -> None
"""Initializes a _GraphGrowingHookDecorator instance.
Args:
hook: The SessionRunHook to decorate.
"""
self._hook = hook
def begin(self):
self._hook.begin()
def end(self, session):
self._hook.end(session)
def _delete_directory(directory):
# type: (Text) -> None
"""Removes directory and handles any folder or file exceptions."""
if not tf.io.gfile.exists(directory):
return
try:
tf.io.gfile.rmtree(directory)
except (tf.errors.PermissionDeniedError,
tf.errors.FailedPreconditionError) as e:
logging.info("Ignoring folder or file issues: %s '%s'", e.error_code,
e.message)
@contextlib.contextmanager
def _disable_asserts_for_confusion_matrix_at_thresholds():
"""Disables asserts in metrics_impl._confusion_matrix_at_thresholds.
AdaNet sometimes have a few NaN and non-NaN subnetworks at a given iteration.
This doesn't crash during training, since AdaNet simply chooses the best
subnetwork among the non-NaN candidates. However, during estimator.evaluate(),
AdaNet evaluates all subnetworks and ensembles. This triggers an assertion
failure in V1 binary classifier _Head, since it expects the predictions to be
between 0 and 1, and NaN is not between 0 and 1. This causes AdaNet and
to raise an exception during estimator.evaluate(), even though the final model
is servable. Hence, we disable these assertions during evaluate(), and allow
the NaNs to be written to disk.
Yields:
Nothing. Simply returns control back to the caller.
"""
from tensorflow.python.ops import metrics_impl # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
def _no_op_assert(x, y, data=None, summarize=None, message=None, name=None):
"""Dummy assert that never fails."""
del x, y, data, summarize, message, name # unused
return tf.no_op()
old_confusion_matrix_at_thresholds = (
metrics_impl._confusion_matrix_at_thresholds) # pylint:disable=protected-access
def _confusion_matrix_at_thresholds_without_asserts(labels,
predictions,
thresholds,
weights=None,
includes=None):
"""Calls _confusion_matrix_at_thresholds without asserts; returns output."""
from tensorflow.python.ops import check_ops # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
old_assert_greater_equal = check_ops.assert_greater_equal
old_assert_less_equal = check_ops.assert_less_equal
setattr(check_ops, "assert_greater_equal", _no_op_assert)
setattr(check_ops, "assert_less_equal", _no_op_assert)
conf_matrix = old_confusion_matrix_at_thresholds(labels, predictions,
thresholds, weights,
includes)
setattr(check_ops, "assert_greater_equal", old_assert_greater_equal)
setattr(check_ops, "assert_less_equal", old_assert_less_equal)
return conf_matrix
setattr(metrics_impl, "_confusion_matrix_at_thresholds",
_confusion_matrix_at_thresholds_without_asserts)
try:
yield
finally:
setattr(metrics_impl, "_confusion_matrix_at_thresholds",
old_confusion_matrix_at_thresholds)
[docs]class Estimator(tf.estimator.Estimator):
# pyformat: disable
r"""A :class:`tf.estimator.Estimator` for training, evaluation, and serving.
This implementation uses an :class:`adanet.subnetwork.Generator` as its weak
learning algorithm for generating candidate subnetworks. These are trained in
parallel using a single graph per iteration. At the end of each iteration, the
estimator saves the sub-graph of the best subnetwork ensemble and its weights
as a separate checkpoint. At the beginning of the next iteration, the
estimator imports the previous iteration's frozen graph and adds ops for the
next candidates as part of a new graph and session. This allows the estimator
have the performance of Tensorflow's static graph constraint (minus the
performance hit of reconstructing a graph between iterations), while having
the flexibility of having a dynamic graph.
NOTE: Subclassing :class:`tf.estimator.Estimator` is only necessary to work
with :meth:`tf.estimator.train_and_evaluate` which asserts that the estimator
argument is a :class:`tf.estimator.Estimator` subclass. However, all training
is delegated to a separate :class:`tf.estimator.Estimator` instance. It is
responsible for supporting both local and distributed training. As such, the
:class:`adanet.Estimator` is only responsible for bookkeeping across
iterations.
Args:
head: A :class:`tf.contrib.estimator.Head` instance for computing loss and
evaluation metrics for every candidate.
subnetwork_generator: The :class:`adanet.subnetwork.Generator` which defines
the candidate subnetworks to train and evaluate at every AdaNet iteration.
max_iteration_steps: Total number of steps for which to train candidates per
iteration. If :class:`OutOfRange` or :class:`StopIteration` occurs in the
middle, training stops before `max_iteration_steps` steps. When
:code:`None`, it will train the current iteration forever.
ensemblers: An iterable of :class:`adanet.ensemble.Ensembler` objects that
define how to ensemble a group of subnetworks. If there are multiple,
each should have a different `name` property.
ensemble_strategies: An iterable of :class:`adanet.ensemble.Strategy`
objects that define the candidate ensembles of subnetworks to explore at
each iteration.
evaluator: An :class:`adanet.Evaluator` for candidate selection after all
subnetworks are done training. When :code:`None`, candidate selection uses
a moving average of their :class:`adanet.Ensemble` AdaNet loss during
training instead. In order to use the *AdaNet algorithm* as described in
[Cortes et al., '17], the given :class:`adanet.Evaluator` must be created
with the same dataset partition used during training. Otherwise, this
framework will perform *AdaNet.HoldOut* which uses a holdout set for
candidate selection, but does not benefit from learning guarantees.
report_materializer: An :class:`adanet.ReportMaterializer`. Its reports are
made available to the `subnetwork_generator` at the next iteration, so
that it can adapt its search space. When `None`, the
`subnetwork_generator` :meth:`generate_candidates` method will receive
empty Lists for their `previous_ensemble_reports` and `all_reports`
arguments.
metric_fn: A function for adding custom evaluation metrics, which should
obey the following signature:
- `Args`:
Can only have the following three arguments in any order:
- :code:`predictions`: Predictions `Tensor` or dict of `Tensor`
created by given :code:`head`.
- :code:`features`: Input `dict` of `Tensor` objects created by
:code:`input_fn` which is given to :meth:`estimator.evaluate` as an
argument.
- :code:`labels`: Labels `Tensor` or dict of `Tensor` (for multi-head)
created by :code:`input_fn` which is given to
:meth:`estimator.evaluate` as an argument.
- `Returns`: Dict of metric results keyed by name. Final metrics are a
union of this and :code:`head`'s existing metrics. If there is a name
conflict between this and :code:`head`s existing metrics, this will
override the existing one. The values of the dict are the results of
calling a metric function, namely a :code:`(metric_tensor, update_op)`
tuple.
force_grow: Boolean override that forces the ensemble to grow by one
subnetwork at the end of each iteration. Normally at the end of each
iteration, AdaNet selects the best candidate ensemble according to its
performance on the AdaNet objective. In some cases, the best ensemble is
the `previous_ensemble` as opposed to one that includes a newly trained
subnetwork. When `True`, the algorithm will not select the
`previous_ensemble` as the best candidate, and will ensure that after n
iterations the final ensemble is composed of n subnetworks.
replicate_ensemble_in_training: Whether to rebuild the frozen subnetworks of
the ensemble in training mode, which can change the outputs of the frozen
subnetworks in the ensemble. When `False` and during candidate training,
the frozen subnetworks in the ensemble are in prediction mode, so
training-only ops like dropout are not applied to them. When `True` and
training the candidates, the frozen subnetworks will be in training mode
as well, so they will apply training-only ops like dropout. This argument
is useful for regularizing learning mixture weights, or for making
training-only side inputs available in subsequent iterations. For most
use-cases, this should be `False`.
adanet_loss_decay: Float decay for the exponential-moving-average of the
AdaNet objective throughout training. This moving average is a data-
driven way tracking the best candidate with only the training set.
delay_secs_per_worker: Float number of seconds to delay starting the
i-th worker. Staggering worker start-up during distributed asynchronous
SGD can improve training stability and speed up convergence. Each worker
will wait (i+1) * delay_secs_per_worker seconds before beginning training.
max_worker_delay_secs: Float max number of seconds to delay starting the
i-th worker. Staggering worker start-up during distributed asynchronous
SGD can improve training stability and speed up convergence. Each worker
will wait up to max_worker_delay_secs before beginning training.
worker_wait_secs: Float number of seconds for workers to wait before
checking if the chief prepared the next iteration.
worker_wait_timeout_secs: Float number of seconds for workers to wait for
chief to prepare the next iteration during distributed training. This is
needed to prevent workers waiting indefinitely for a chief that may have
crashed or been turned down. When the timeout is exceeded, the worker
exits the train loop. In situations where the chief job is much slower
than the worker jobs, this timeout should be increased.
model_dir: Directory to save model parameters, graph and etc. This can also
be used to load checkpoints from the directory into a estimator to
continue training a previously saved model.
report_dir: Directory where the
:class:`adanet.subnetwork.MaterializedReport`s materialized by
:code:`report_materializer` would be saved. If :code:`report_materializer`
is :code:`None`, this will not save anything. If :code:`None` or
empty string, defaults to :code:`<model_dir>/report`.
config: :class:`RunConfig` object to configure the runtime settings.
debug: Boolean to enable debug mode which will check features and labels
for Infs and NaNs.
enable_ensemble_summaries: Whether to record summaries to display in
TensorBoard for each ensemble candidate. Disable to reduce memory and disk
usage per run.
enable_subnetwork_summaries: Whether to record summaries to display in
TensorBoard for each subnetwork. Disable to reduce memory and disk usage
per run.
global_step_combiner_fn: Function for combining each subnetwork's
iteration step into the global step. By default it is the average of all
subnetwork iteration steps, which may affect the global_steps/sec as
subnetworks early stop and no longer increase their iteration step.
max_iterations: Integer maximum number of AdaNet iterations (a.k.a. rounds)
of generating new subnetworks and ensembles, training them, and evaluating
them against the current best ensemble. When :code:`None`, AdaNet will
keep iterating until `Estimator#train` terminates. Otherwise, if
:code:`max_iteratios` is supplied and is met or exceeded during training,
training will terminate even before `steps` or `max_steps`.
export_subnetwork_logits: Whether to include subnetwork logits in exports.
export_subnetwork_last_layer: Whether to include subnetwork last layer in
exports.
replay_config: Optional :class:`adanet.replay.Config` to specify a previous
AdaNet run to replay. Given the exact same search space but potentially
different training data, the `replay_config` causes the estimator to
reconstruct the previously trained model without performing a search.
NOTE: The previous run must have executed with identical hyperparameters
as the new run in order to be replayable. The only supported difference is
that the underlying data can change.
**kwargs: Extra keyword args passed to the parent.
Returns:
An :class:`adanet.Estimator` instance.
Raises:
:code:`ValueError`: If :code:`subnetwork_generator` is :code:`None`.
:code:`ValueError`: If :code:`max_iteration_steps` is <= 0.
:code:`ValueError`: If :code:`model_dir` is not specified during distributed
training.
:code:`ValueError`: If :code:`max_iterations` is <= 0.
"""
# pyformat: enable
class _Keys(object):
CURRENT_ITERATION = "current_iteration"
SUBNETWORK_GENERATOR = "subnetwork_generator"
def __init__(self,
head,
subnetwork_generator,
max_iteration_steps,
ensemblers=None,
ensemble_strategies=None,
evaluator=None,
report_materializer=None,
metric_fn=None,
force_grow=False,
replicate_ensemble_in_training=False,
adanet_loss_decay=.9,
delay_secs_per_worker=5,
max_worker_delay_secs=60,
worker_wait_secs=5,
worker_wait_timeout_secs=7200,
model_dir=None,
report_dir=None,
config=None,
debug=False,
enable_ensemble_summaries=True,
enable_subnetwork_summaries=True,
global_step_combiner_fn=tf.math.reduce_mean,
max_iterations=None,
export_subnetwork_logits=False,
export_subnetwork_last_layer=True,
replay_config=None,
**kwargs):
if subnetwork_generator is None:
raise ValueError("subnetwork_generator can't be None.")
if max_iteration_steps is not None and max_iteration_steps <= 0.:
raise ValueError("max_iteration_steps must be > 0 or None.")
if max_iterations is not None and max_iterations <= 0.:
raise ValueError("max_iterations must be > 0 or None.")
is_distributed_training = config and config.num_worker_replicas > 1
is_model_dir_specified = model_dir or (config and config.model_dir)
if is_distributed_training and not is_model_dir_specified:
# A common model dir for the chief and workers is required for
# coordination during distributed training.
raise ValueError(
"For distributed training, a model_dir must be specified.")
self._subnetwork_generator = subnetwork_generator
# Overwrite superclass's assert that members are not overwritten in order
# to overwrite public methods. Note that we are doing something that is not
# explicitly supported by the Estimator API and may break in the future.
tf.estimator.Estimator._assert_members_are_not_overridden = staticmethod( # pylint: disable=protected-access
lambda _: None)
self._enable_v2_checkpoint = kwargs.pop("enable_v2_checkpoint", False)
self._evaluator = evaluator
self._report_materializer = report_materializer
self._force_grow = force_grow
self._delay_secs_per_worker = delay_secs_per_worker
self._max_worker_delay_secs = max_worker_delay_secs
self._worker_wait_secs = worker_wait_secs
self._worker_wait_timeout_secs = worker_wait_timeout_secs
self._max_iterations = max_iterations
self._replay_config = replay_config
# Added for backwards compatibility.
default_ensembler_args = [
"mixture_weight_type", "mixture_weight_initializer",
"warm_start_mixture_weights", "adanet_lambda", "adanet_beta", "use_bias"
]
default_ensembler_kwargs = {
k: v for k, v in kwargs.items() if k in default_ensembler_args
}
if default_ensembler_kwargs:
logging.warning(
"The following arguments have been moved to "
"`adanet.ensemble.ComplexityRegularizedEnsembler` which can be "
"specified in the `ensemblers` argument: %s",
sorted(default_ensembler_kwargs.keys()))
for key in default_ensembler_kwargs:
del kwargs[key]
# Experimental feature.
placement_strategy_arg = "experimental_placement_strategy"
placement_strategy = kwargs.pop(placement_strategy_arg, None)
if placement_strategy:
logging.warning(
"%s is an experimental feature. Its behavior is not guaranteed "
"to be backwards compatible.", placement_strategy_arg)
self._warm_start_settings = kwargs.get("warm_start_from")
# Monkey patch the default variable placement strategy that Estimator uses
# since it does not support workers having different graphs from the chief.
# TODO: Consider using `RunConfig.replace` with the new device_fn,
# but this can cause issues since RunConfig automatically parses TF_CONFIG
# environment variable.
with monkey_patch_default_variable_placement_strategy():
# This `Estimator` is responsible for bookkeeping across iterations, and
# for training the subnetworks in both a local and distributed setting.
# Subclassing improves future-proofing against new private methods being
# added to `tf.estimator.Estimator` that are expected to be callable by
# external functions, such as in b/110435640.
super(Estimator, self).__init__(
model_fn=self._create_model_fn(),
params={},
config=config,
model_dir=model_dir,
**kwargs)
if default_ensembler_kwargs and ensemblers:
raise ValueError("When specifying the `ensemblers` argument, "
"the following arguments must not be given: {}".format(
default_ensembler_kwargs.keys()))
if not ensemblers:
default_ensembler_kwargs["model_dir"] = self.model_dir
ensemblers = [
ensemble_lib.ComplexityRegularizedEnsembler(
**default_ensembler_kwargs)
]
# These are defined after base Estimator's init so that they can
# use the same temporary model_dir as the underlying Estimator even if
# model_dir is not provided.
self._use_tpu = kwargs.get("use_tpu", False)
ensemble_builder = _EnsembleBuilder(
head=head,
metric_fn=metric_fn,
use_tpu=self._use_tpu,
export_subnetwork_logits=export_subnetwork_logits,
export_subnetwork_last_layer=export_subnetwork_last_layer)
# TODO: Merge CandidateBuilder into SubnetworkManager.
candidate_builder = _CandidateBuilder(adanet_loss_decay=adanet_loss_decay)
subnetwork_manager = _SubnetworkManager(
head=head, metric_fn=metric_fn, use_tpu=self._use_tpu)
if not placement_strategy:
placement_strategy = distributed_lib.ReplicationStrategy()
self._iteration_builder = _IterationBuilder(
candidate_builder,
subnetwork_manager,
ensemble_builder,
ensemblers,
max_iteration_steps,
self._summary_maker,
global_step_combiner_fn,
placement_strategy,
replicate_ensemble_in_training,
use_tpu=self._use_tpu,
debug=debug,
enable_ensemble_summaries=enable_ensemble_summaries,
enable_subnetwork_summaries=enable_subnetwork_summaries,
enable_subnetwork_reports=self._report_materializer is not None)
self._ensemble_strategies = ensemble_strategies or [
ensemble_lib.GrowStrategy()
]
report_dir = report_dir or os.path.join(self._model_dir, "report")
self._report_accessor = _ReportAccessor(report_dir)
def _summary_maker(self, scope=None, skip_summary=False, namespace=None):
"""Constructs a `_ScopedSummary`."""
if tf_compat.is_v2_behavior_enabled():
# Here we assume TF 2 behavior is enabled.
return _ScopedSummaryV2(
logdir=self._model_dir,
scope=scope,
skip_summary=skip_summary,
namespace=namespace)
if self._use_tpu:
return _TPUScopedSummary(
logdir=self._model_dir,
scope=scope,
skip_summary=skip_summary,
namespace=namespace)
else:
return _ScopedSummary(
scope=scope, skip_summary=skip_summary, namespace=namespace)
def _checkpoint_iteration_number(self, checkpoint_path):
# type: (Text) -> int
"""Returns the iteration number from the latest checkpoint."""
if checkpoint_path is None:
return 0
if self._enable_v2_checkpoint:
return tf_compat.load_variable(
checkpoint_path, "iteration_number", shape=[], dtype=tf.int64)
return tf.train.load_variable(checkpoint_path,
self._Keys.CURRENT_ITERATION).item()
def _checkpoint_global_step(self, checkpoint_path):
# type: (Text) -> int
"""Returns the global step from the given checkpoint."""
if checkpoint_path is None:
return 0
if self._enable_v2_checkpoint:
return tf_compat.load_variable(
checkpoint_path,
tf_compat.v1.GraphKeys.GLOBAL_STEP,
shape=[],
dtype=tf.int64)
return tf.train.load_variable(checkpoint_path,
tf_compat.v1.GraphKeys.GLOBAL_STEP).item()
[docs] def train(self,
input_fn,
hooks=None,
steps=None,
max_steps=None,
saving_listeners=None):
# pyformat: disable
"""Trains a model given training data :code:`input_fn`.
NOTE: If a given input_fn raises an :code:`OutOfRangeError`, then *all* of
training will exit. The best practice is to make the training dataset repeat
forever, in order to perform model search for more than one iteration.
Args:
input_fn: A function that provides input data for training as minibatches.
See [Premade Estimators](
https://tensorflow.org/guide/premade_estimators#create_input_functions)
for more information. The function should construct and return one of
the following:
* A :code:`tf.data.Dataset` object: Outputs of `Dataset` object must
be a tuple `(features, labels)` with same constraints as below.
* A tuple `(features, labels)`: Where `features` is a
:code:`tf.Tensor` or a dictionary of string feature name to
`Tensor` and `labels` is a :code:`Tensor` or a dictionary of string
label name to `Tensor`. Both `features` and `labels` are consumed by
`model_fn`. They should satisfy the expectation of `model_fn` from
inputs.
hooks: List of :code:`tf.train.SessionRunHook` subclass instances. Used
for callbacks inside the training loop.
steps: Number of steps for which to train the model. If :code:`None`,
train forever or train until `input_fn` generates the
:code:`tf.errors.OutOfRange` error or :code:`StopIteration` exception.
`steps` works incrementally. If you call two times `train(steps=10)`
then training occurs in total 20 steps. If :code:`OutOfRange` or
:code:`StopIteration` occurs in the middle, training stops before 20
steps. If you don't want to have incremental behavior please set
`max_steps` instead. If set, `max_steps` must be :code:`None`.
max_steps: Number of total steps for which to train model. If
:code:`None`, train forever or train until `input_fn` generates the
:code:`tf.errors.OutOfRange` error or :code:`StopIteration` exception.
If set, `steps` must be `None`. If :code:`OutOfRange` or
:code:`StopIteration` occurs in the middle, training stops before
`max_steps` steps. Two calls to `train(steps=100)` means 200 training
iterations. On the other hand, two calls to `train(max_steps=100)`
means that the second call will not do any iteration since first call
did all 100 steps.
saving_listeners: list of :code:`CheckpointSaverListener` objects. Used
for callbacks that run immediately before or after checkpoint savings.
Returns:
`self`, for chaining.
Raises:
ValueError: If both `steps` and `max_steps` are not `None`.
ValueError: If either `steps` or `max_steps <= 0`.
"""
# pyformat: enable
if (steps is not None) and (max_steps is not None):
raise ValueError("Can not provide both steps and max_steps.")
if steps is not None and steps <= 0:
raise ValueError("Must specify steps > 0, given: {}".format(steps))
latest_checkpoint = tf.train.latest_checkpoint(self.model_dir)
latest_global_steps = self._checkpoint_global_step(latest_checkpoint)
if steps is not None:
max_steps = latest_global_steps + steps
# Each iteration of this AdaNet loop represents an `_Iteration`. The
# current iteration number is stored as a variable in the checkpoint so
# that training can be stopped and started at anytime.
with monkey_patch_default_variable_placement_strategy():
while True:
latest_checkpoint = tf.train.latest_checkpoint(self.model_dir)
latest_global_steps = self._checkpoint_global_step(latest_checkpoint)
current_iteration = self._checkpoint_iteration_number(latest_checkpoint)
logging.info("Beginning training AdaNet iteration %s",
current_iteration)
self._iteration_ended = False
# Delegate training to a temporary estimator instead of super to make
# passing arguments more functional (via params).
temp_estimator = self._create_temp_estimator(
config=self.config,
is_inside_training_loop=True,
checkpoint_path=latest_checkpoint,
hooks=hooks)
result = temp_estimator.train(
input_fn=input_fn,
hooks=hooks,
max_steps=max_steps,
saving_listeners=saving_listeners)
# In TensorFlow v2.0.0.rc1 and below, saving listeners are attached to
# the first CheckpointSaverHook each time train is called. Instead, we
# pass in the saving_listeners in the first AdaNet iteration only.
if not tf_compat.version_greater_or_equal("2.0.0.rc1"):
saving_listeners = None
logging.info("Finished training Adanet iteration %s", current_iteration)
# If training ended because the maximum number of training steps
# occurred, exit training.
latest_checkpoint = tf.train.latest_checkpoint(self.model_dir)
global_steps = self._checkpoint_global_step(latest_checkpoint)
if max_steps is not None and global_steps >= max_steps:
logging.info("Training ended after %s global steps", global_steps)
return result
# If training ended for any reason other than the iteration ending,
# exit training.
if not self._iteration_ended:
logging.info("Training stop requested")
return result
max_iterations = self._max_iterations
if max_iterations and current_iteration + 1 >= max_iterations:
logging.info(
"Training ended after exceeding maximum AdaNet iterations")
if steps is not None and global_steps - latest_global_steps < steps:
logging.warning(
"Both `max_iterations` and `steps` were specified, but "
"`max_iterations` takes precedence over `steps`")
return result
logging.info("Beginning bookkeeping phase for iteration %s",
current_iteration)
# The chief prepares the next AdaNet iteration, and increments the
# iteration number by 1.
if self.config.is_chief:
with self._force_replication_strategy():
self._execute_bookkeeping_phase(
input_fn,
current_iteration,
train_hooks=hooks or [],
checkpoint_path=latest_checkpoint)
# This inner loop serves mainly for synchronizing the workers with the
# chief during distributed training. Workers that finish training early
# wait for the chief to prepare the next iteration and increment the
# iteration number. Workers that are slow to finish training quickly
# move onto the next iteration. And workers that go offline and return
# online after training ended terminate gracefully.
wait_for_chief = not self.config.is_chief
timer = _CountDownTimer(self._worker_wait_timeout_secs)
while wait_for_chief:
# Fetch the latest checkpoint.
latest_checkpoint = tf.train.latest_checkpoint(self.model_dir)
# If the chief hits max_steps, it will stop training itself and not
# increment the iteration number, so this is how the worker knows to
# exit if it wakes up and the chief is gone.
# TODO: Support steps parameter.
if self._checkpoint_global_step(latest_checkpoint) >= max_steps:
return result
# In distributed training, a worker may end training before the chief
# overwrites the checkpoint with the incremented iteration number. If
# that is the case, it should wait for the chief to do so. Otherwise
# the worker will get stuck waiting for its weights to be initialized.
next_iteration = self._checkpoint_iteration_number(latest_checkpoint)
if next_iteration > current_iteration:
break
logging.info("Iteration number in latest checkpoint: %d",
next_iteration)
# Check timeout when waiting for potentially downed chief.
if timer.secs_remaining() == 0:
logging.error(
"Chief job did not prepare iteration %d after %s secs. It "
"may have been preempted, been turned down, or crashed. This "
"worker is now exiting training.", current_iteration + 1,
self._worker_wait_timeout_secs)
return result
logging.info("Waiting for chief to prepare iteration %d",
current_iteration + 1)
time.sleep(self._worker_wait_secs)
# Stagger starting workers to prevent training instability.
# Mimics behavior of tf.estimator.train_and_evaluate.
if not self.config.is_chief and self.config.task_type == "worker":
task_id = self.config.task_id or 0
# Stagger each worker up to 60 secs.
delay_secs = min(self._max_worker_delay_secs,
(task_id + 1.) * self._delay_secs_per_worker)
if delay_secs > 0.:
logging.info("Waiting %d secs before continuing training.",
delay_secs)
time.sleep(delay_secs)
logging.info("Finished bookkeeping phase for iteration %s",
current_iteration)
[docs] def evaluate(self,
input_fn,
steps=None,
hooks=None,
checkpoint_path=None,
name=None):
if not checkpoint_path:
checkpoint_path = tf.train.latest_checkpoint(self.model_dir)
logging.info("Evaluating AdaNet model at checkpoint: %s", checkpoint_path)
# Delegate evaluation to a temporary estimator instead of super to make
# passing arguments more functional (via params).
temp_estimator = self._create_temp_estimator(
config=self.config,
checkpoint_path=checkpoint_path,
evaluation_name=name,
# Ensure that the read to get the iteration number and read to restore
# variable values come from the same checkpoint during evaluation.
best_ensemble_index=self._compute_best_ensemble_index(
checkpoint_path, hooks),
hooks=hooks)
with _disable_asserts_for_confusion_matrix_at_thresholds():
result = temp_estimator.evaluate(
input_fn,
steps=steps,
hooks=hooks,
checkpoint_path=checkpoint_path,
name=name)
return result
[docs] def predict(self,
input_fn,
predict_keys=None,
hooks=None,
checkpoint_path=None,
yield_single_examples=True):
if not checkpoint_path:
checkpoint_path = tf.train.latest_checkpoint(self.model_dir)
logging.info("Computing predictions for AdaNet model at checkpoint: %s",
checkpoint_path)
# Delegate predicting to a temporary estimator instead of super to make
# passing arguments more functional (via params).
temp_estimator = self._create_temp_estimator(
config=self.config,
best_ensemble_index=self._compute_best_ensemble_index(
checkpoint_path, hooks=hooks),
checkpoint_path=checkpoint_path,
hooks=hooks)
return temp_estimator.predict(
input_fn=input_fn,
predict_keys=predict_keys,
checkpoint_path=checkpoint_path,
yield_single_examples=yield_single_examples,
hooks=hooks)
from tensorflow.python.util import deprecation # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
[docs] @deprecation.deprecated(
None, "This function has been renamed, use `export_saved_model` instead.")
def export_savedmodel(self,
export_dir_base,
serving_input_receiver_fn,
hooks=None,
assets_extra=None,
as_text=False,
checkpoint_path=None,
strip_default_attrs=False):
if not checkpoint_path:
checkpoint_path = tf.train.latest_checkpoint(self.model_dir)
logging.info("Exporting SavedModel for AdaNet model at checkpoint: %s",
checkpoint_path)
# Delegate exporting to a temporary estimator instead of super to make
# passing arguments more functional (via params).
temp_estimator = self._create_temp_estimator(
config=self.config,
hooks=hooks,
best_ensemble_index=self._compute_best_ensemble_index(
checkpoint_path, hooks=hooks),
checkpoint_path=checkpoint_path,
is_export=True)
with self._force_replication_strategy():
return temp_estimator.export_savedmodel(
export_dir_base=export_dir_base,
serving_input_receiver_fn=serving_input_receiver_fn,
assets_extra=assets_extra,
as_text=as_text,
checkpoint_path=checkpoint_path,
strip_default_attrs=strip_default_attrs)
[docs] def export_saved_model(self,
export_dir_base,
serving_input_receiver_fn,
hooks=None,
assets_extra=None,
as_text=False,
checkpoint_path=None,
experimental_mode=tf.estimator.ModeKeys.PREDICT):
if not checkpoint_path:
checkpoint_path = tf.train.latest_checkpoint(self.model_dir)
logging.info("Exporting SavedModel for AdaNet model at checkpoint: %s",
checkpoint_path)
# Delegate exporting to a temporary estimator instead of super to make
# passing arguments more functional (via params).
temp_estimator = self._create_temp_estimator(
config=self.config,
best_ensemble_index=self._compute_best_ensemble_index(
checkpoint_path, hooks=hooks),
checkpoint_path=checkpoint_path,
hooks=hooks,
is_export=True)
with self._force_replication_strategy():
return temp_estimator.export_saved_model(
export_dir_base=export_dir_base,
serving_input_receiver_fn=serving_input_receiver_fn,
assets_extra=assets_extra,
as_text=as_text,
checkpoint_path=checkpoint_path,
experimental_mode=experimental_mode)
[docs] def experimental_export_all_saved_models(self,
export_dir_base,
input_receiver_fn_map,
hooks=None,
assets_extra=None,
as_text=False,
checkpoint_path=None):
if not checkpoint_path:
checkpoint_path = tf.train.latest_checkpoint(self.model_dir)
logging.info("Exporting SavedModel for AdaNet model at checkpoint: %s",
checkpoint_path)
# Delegate exporting to a temporary estimator instead of super to make
# passing arguments more functional (via params).
temp_estimator = self._create_temp_estimator(
config=self.config,
best_ensemble_index=self._compute_best_ensemble_index(
checkpoint_path, hooks=hooks),
checkpoint_path=checkpoint_path,
hooks=hooks,
is_export=True)
with self._force_replication_strategy():
return temp_estimator.experimental_export_all_saved_models(
export_dir_base=export_dir_base,
input_receiver_fn_map=input_receiver_fn_map,
assets_extra=assets_extra,
as_text=as_text,
checkpoint_path=checkpoint_path)
def _compute_best_ensemble_index(self, checkpoint_path, hooks):
# type: (Text, Sequence[tf_compat.SessionRunHook]) -> Optional[int]
"""Runs the Evaluator to obtain the best ensemble index among candidates."""
# AdaNet Replay.
if self._replay_config:
iteration_number = self._checkpoint_iteration_number(checkpoint_path)
best_index = self._replay_config.get_best_ensemble_index(iteration_number)
if best_index is not None:
return best_index
if self._evaluator:
return self._execute_candidate_evaluation_phase(
self._evaluator.input_fn,
export_best_architecture=False,
checkpoint_path=checkpoint_path,
hooks=hooks)
return None
@contextlib.contextmanager
def _force_replication_strategy(self):
"""Sets placement_strategy to always be ReplicationStrategy.
This is useful during the bookkeeping phase and when Estimator's export
saved model functions are called. In both of these cases, local tf.Sessions
are created which do not have access to the cluster. Therefore,
RoundRobinReplicationStrategy will fail when it tries to place ops on
cluster devices which the local tf.Sessions cannot access.
Yields:
Nothing. Simply returns control back to the caller.
"""
temp_placement_strategy = self._iteration_builder.placement_strategy
try:
placement_strategy = distributed_lib.ReplicationStrategy()
self._iteration_builder.placement_strategy = placement_strategy
yield
finally:
self._iteration_builder.placement_strategy = temp_placement_strategy
@contextlib.contextmanager
def _call_input_fn_in_new_graph(self, input_fn, mode, config):
"""Calls the given input_fn and yields results within a new graph context.
Yields features, labels, and hooks from the result of an Estimator input_fn.
Args:
input_fn: a function that takes no arguments and returns one of:
* A 'tf.data.Dataset' object: Outputs of `Dataset` object must be a
tuple (features, labels) with same constraints as below.
* A tuple (features, labels): Where `features` is a `Tensor` or a
dictionary of string feature name to `Tensor` and `labels` is a
`Tensor` or a dictionary of string label name to `Tensor`. Both
`features` and `labels` are consumed by `model_fn`. They should
satisfy the expectation of `model_fn` from inputs.
mode: Defines whether this is training, evaluation or prediction. See
`ModeKeys`.
config: The current `tf.estimator.RunConfig`.
Yields:
Tuple of features, labels, and input_hooks, where features are as
described above, labels are as described above or None, and input_hooks
are a list of SessionRunHooks to be included when running.
Raises:
ValueError: if the result is a list or tuple of length != 2.
"""
from tensorflow_estimator.python.estimator import util # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
with tf.Graph().as_default() as g:
tf_compat.v1.set_random_seed(config.tf_random_seed)
# Create global step before calling model_fn as does superclass.
self._create_and_assert_global_step(g)
with tf.device("/cpu:0"):
input_fn_outs = input_fn()
yield util.parse_input_fn_result(input_fn_outs)
def _create_temp_run_config(self, temp_model_dir):
# type: (Text) -> tf.estimator.RunConfig
"""Creates a temp `RunConfig` for the bookkeeping phase."""
config = self.config
return tf.estimator.RunConfig(
model_dir=temp_model_dir,
tf_random_seed=config.tf_random_seed,
session_config=config.session_config,
protocol=config.protocol)
def _create_temp_estimator(self, config, **create_model_fn_args):
# type: (tf.estimator.RunConfig, Any[...]) -> tf.estimator.Estimator # pylint:disable=line-too-long
"""Creates a temp `Estimator` to grow the graph for the next iteration."""
return tf.estimator.Estimator(
model_fn=self._create_model_fn(**create_model_fn_args),
config=config,
warm_start_from=self._warm_start_settings)
def _execute_bookkeeping_phase(self, train_input_fn, iteration_number,
train_hooks, checkpoint_path):
"""Run the AdaNet bookkeeping phase to prepare the next iteration.
This method creates a TensorFlow graph up to three times:
1. To evaluate all candidate ensembles to find the best one.
2. To materialize reports and store them to disk (if report_materializer
exists).
3. To grow the TensorFlow graph and overwrite the model directory's
checkpoint with the next iteration's ops.
Args:
train_input_fn: The input_fn used during training.
iteration_number: Integer current iteration number.
train_hooks: List of `SessionRunHook` passed for training.
checkpoint_path: Path to the checkpoint to restore from.
"""
next_iteration_number = iteration_number + 1
logging.info("Preparing iteration %s:", next_iteration_number)
if self._evaluator:
evaluator_input_fn = self._evaluator.input_fn
else:
evaluator_input_fn = train_input_fn
best_ensemble_index = self._execute_candidate_evaluation_phase(
evaluator_input_fn,
export_best_architecture=True,
checkpoint_path=checkpoint_path,
hooks=train_hooks)
self._execute_report_materialization_phase(
best_ensemble_index, checkpoint_path=checkpoint_path, hooks=train_hooks)
self._execute_graph_growing_phase(train_input_fn, train_hooks,
checkpoint_path)
logging.info("Finished preparing iteration %s.", next_iteration_number)
def _execute_candidate_evaluation_phase(self, evaluator_input_fn,
export_best_architecture,
checkpoint_path, hooks):
"""Evaluates and chooses the best ensemble for this iteration.
Args:
evaluator_input_fn: The input_fn for evaluation.
export_best_architecture: Boolean whether to persist the best ensemble's
architecture to the model_dir.
checkpoint_path: Path to the checkpoint to restore from.
hooks: A list of `tf.estimator.SessionRunHook`s.
Returns:
Integer index of the best ensemble withing the list of candidate ensembles
for the current iteration.
"""
logging.info("Evaluating candidates...")
config = self.config
mode = tf.estimator.ModeKeys.EVAL
with self._call_input_fn_in_new_graph(evaluator_input_fn, mode,
config) as (features, labels,
input_hooks):
current_iteration, _ = self._create_iteration(
features,
labels,
mode,
config,
is_growing_phase=False,
checkpoint_path=checkpoint_path,
hooks=hooks)
best_ensemble_index = self._get_best_ensemble_index(
current_iteration, input_hooks, checkpoint_path)
architecture = current_iteration.candidates[
best_ensemble_index].ensemble_spec.architecture
if export_best_architecture:
iteration_number = self._checkpoint_iteration_number(checkpoint_path)
new_architecture_filename = self._architecture_filename(iteration_number)
logging.info("Exporting best ensemble architecture to %s",
new_architecture_filename)
self._save_architecture(new_architecture_filename, architecture,
checkpoint_path)
logging.info("Done evaluating candidates.")
return best_ensemble_index
def _execute_report_materialization_phase(self, best_ensemble_index,
checkpoint_path, hooks):
"""Materializes and store subnetwork reports."""
if not self._report_materializer:
return
logging.info("Materializing reports...")
input_fn = self._report_materializer.input_fn
mode = tf.estimator.ModeKeys.EVAL
config = self.config
with self._call_input_fn_in_new_graph(input_fn, mode,
config) as (features, labels,
input_hooks):
current_iteration, _ = self._create_iteration(
features,
labels,
mode,
config,
is_growing_phase=False,
checkpoint_path=checkpoint_path,
hooks=hooks)
self._materialize_report(current_iteration, input_hooks,
best_ensemble_index, checkpoint_path)
logging.info("Done materializing reports.")
def _execute_graph_growing_phase(self, train_input_fn, train_hooks,
checkpoint_path):
"""Grows the tensorflow graph for the next iteration.
Normally the MonitoredTrainingSession does not allow one to add new ops to
the TensorFlow graph once training starts. To get around this limitation,
create the graph for the next iteration and overwrite the model directory
checkpoint with the expanded graph.
Args:
train_input_fn: The input_fn used during training.
train_hooks: List of `SessionRunHook` passed for training.
checkpoint_path: Path of the checkpoint to use for restoring variables.
"""
logging.info("Adapting graph and incrementing iteration number...")
config = self.config
temp_model_dir = os.path.join(self.model_dir, "temp_model_dir")
if not tf.io.gfile.exists(temp_model_dir):
tf.io.gfile.makedirs(temp_model_dir)
# Since deleting a model_dir can fail, we need each temporary directory to
# be unique. So we use the UTC time when creating it.
time_in_millis = int(time.time() * 1000)
temp_model_sub_dir = os.path.join(temp_model_dir, str(time_in_millis))
temp_run_config = config.replace(model_dir=temp_model_sub_dir)
temp_estimator = self._create_temp_estimator(
config=temp_run_config,
is_growing_phase=True,
is_inside_training_loop=True,
checkpoint_path=checkpoint_path,
hooks=train_hooks)
_copy_recursively(
os.path.join(self._model_dir, "assets"),
os.path.join(temp_model_sub_dir, "assets"))
# Do not train with any saving_listeners since this is just a temporary
# estimator.
temp_estimator.train(
input_fn=train_input_fn,
max_steps=1,
hooks=self._process_hooks_for_growing_phase(train_hooks),
saving_listeners=None)
_copy_recursively(
os.path.join(temp_model_sub_dir, "assets"),
os.path.join(self._model_dir, "assets"))
_delete_directory(temp_model_dir)
logging.info("Done adapting graph and incrementing iteration number.")
def _architecture_filename(self, iteration_number):
# type: (int) -> Text
"""Returns the filename of the given iteration's frozen graph."""
frozen_checkpoint = os.path.join(self.model_dir, "architecture")
return "{}-{}.json".format(frozen_checkpoint, iteration_number)
def _get_best_ensemble_index(self,
current_iteration,
input_hooks,
checkpoint_path=None):
# type: (_Iteration, Sequence[tf_compat.SessionRunHook], Text) -> int
"""Returns the best candidate ensemble's index in this iteration.
Evaluates the ensembles using an `Evaluator` when provided. Otherwise,
it returns the index of the best candidate as defined by the `_Iteration`.
Args:
current_iteration: Current `_Iteration`.
input_hooks: List of SessionRunHooks to be included when running.
checkpoint_path: Checkpoint to use when determining the best index.
Returns:
Index of the best ensemble in the iteration's list of `_Candidates`.
"""
# AdaNet Replay.
if self._replay_config:
best_index = self._replay_config.get_best_ensemble_index(
current_iteration.number)
if best_index is not None:
return best_index
# Skip the evaluation phase when there is only one candidate subnetwork.
if len(current_iteration.candidates) == 1:
logging.info("'%s' is the only ensemble",
current_iteration.candidates[0].ensemble_spec.name)
return 0
# The zero-th index candidate at iteration t>0 is always the
# previous_ensemble.
if current_iteration.number > 0 and self._force_grow and (len(
current_iteration.candidates) == 2):
logging.info("With `force_grow` enabled, '%s' is the only ensemble",
current_iteration.candidates[1].ensemble_spec.name)
return 1
logging.info("Starting ensemble evaluation for iteration %s",
current_iteration.number)
for hook in input_hooks:
hook.begin()
with tf_compat.v1.Session(config=self.config.session_config) as sess:
init = tf.group(
tf_compat.v1.global_variables_initializer(),
tf_compat.v1.local_variables_initializer(),
tf_compat.v1.tables_initializer(),
current_iteration.estimator_spec.scaffold.local_init_op if isinstance(
current_iteration.estimator_spec,
tf.estimator.EstimatorSpec) else tf.no_op())
sess.run(init)
if self._enable_v2_checkpoint:
status = current_iteration.checkpoint.restore(checkpoint_path)
status.expect_partial() # Optional sanity checks.
status.initialize_or_restore(sess)
else:
saver = tf_compat.v1.train.Saver(sharded=True)
saver.restore(sess, checkpoint_path)
coord = tf.train.Coordinator()
for hook in input_hooks:
hook.after_create_session(sess, coord)
tf_compat.v1.train.start_queue_runners(sess=sess, coord=coord)
ensemble_metrics = []
for candidate in current_iteration.candidates:
metrics = candidate.ensemble_spec.eval_metrics.eval_metrics_ops()
metrics["adanet_loss"] = tf_compat.v1.metrics.mean(
candidate.ensemble_spec.adanet_loss)
ensemble_metrics.append(metrics)
if self._evaluator:
metric_name = self._evaluator.metric_name
metrics = self._evaluator.evaluate(sess, ensemble_metrics)
objective_fn = self._evaluator.objective_fn
else:
metric_name = "adanet_loss"
metrics = sess.run(
[c.adanet_loss for c in current_iteration.candidates])
objective_fn = np.nanargmin
values = []
for i in range(len(current_iteration.candidates)):
ensemble_name = current_iteration.candidates[i].ensemble_spec.name
values.append("{}/{} = {:.6f}".format(metric_name, ensemble_name,
metrics[i]))
logging.info("Computed ensemble metrics: %s", ", ".join(values))
if self._force_grow and current_iteration.number > 0:
logging.info(
"The `force_grow` override is enabled, so the "
"the performance of the previous ensemble will be ignored.")
# NOTE: The zero-th index candidate at iteration t>0 is always the
# previous_ensemble.
metrics = metrics[1:]
index = objective_fn(metrics) + 1
else:
index = objective_fn(metrics)
logging.info("Finished ensemble evaluation for iteration %s",
current_iteration.number)
logging.info("'%s' at index %s is the best ensemble",
current_iteration.candidates[index].ensemble_spec.name, index)
return index
def _materialize_report(self, current_iteration, input_hooks,
best_ensemble_index, checkpoint_path):
"""Generates reports as defined by `Builder`s.
Materializes the Tensors and metrics defined in the `Builder`s'
`build_subnetwork_report` method using `ReportMaterializer`, and stores
them to disk using `_ReportAccessor`.
Args:
current_iteration: Current `_Iteration`.
input_hooks: List of SessionRunHooks to be included when running.
best_ensemble_index: Integer index of the best candidate ensemble.
checkpoint_path: Path of the checkpoint to use.
"""
logging.info("Starting metric logging for iteration %s",
current_iteration.number)
best_candidate = current_iteration.candidates[best_ensemble_index]
best_architecture = best_candidate.ensemble_spec.architecture
included_subnetwork_names = [
name for i, name in best_architecture.subnetworks
if i == current_iteration.number
]
for hook in input_hooks:
hook.begin()
if self._enable_v2_checkpoint:
status = current_iteration.checkpoint.restore(checkpoint_path)
# Verify that restoring subset of ops from previous iteration works.
status.expect_partial() # Optional sanity checks.
with tf_compat.v1.Session(config=self.config.session_config) as sess:
init = tf.group(
tf_compat.v1.global_variables_initializer(),
tf_compat.v1.local_variables_initializer(),
tf_compat.v1.tables_initializer(),
current_iteration.estimator_spec.scaffold.local_init_op if isinstance(
current_iteration.estimator_spec,
tf.estimator.EstimatorSpec) else tf.no_op())
sess.run(init)
if self._enable_v2_checkpoint:
status.initialize_or_restore(sess)
else:
saver = tf_compat.v1.train.Saver(sharded=True)
saver.restore(sess, checkpoint_path)
coord = tf.train.Coordinator()
for hook in input_hooks:
hook.after_create_session(sess, coord)
tf_compat.v1.train.start_queue_runners(sess=sess, coord=coord)
materialized_reports = (
self._report_materializer.materialize_subnetwork_reports(
sess, current_iteration.number,
current_iteration.subnetwork_reports, included_subnetwork_names))
self._report_accessor.write_iteration_report(current_iteration.number,
materialized_reports)
logging.info("Finished saving subnetwork reports for iteration %s",
current_iteration.number)
def _process_hooks_for_growing_phase(self, hooks):
"""Processes hooks which will run during the graph growing phase.
In particular the following things are done:
- CheckpointSaverHooks are filtered out since they are not intended to
run between training runs and will cause errors. We also reset the
CheckpointSaverHooks' Saver between iterations, see b/122795064 for more
details.
- Decorate the remaining hooks with _GraphGrowingHookDecorator to only run
the begin() and end() methods during the graph growing phase.
Args:
hooks: The list of `SessionRunHooks` to process.
Returns:
The processed hooks which should run during the growing phase.
"""
processed_hooks = []
for hook in hooks:
# Reset CheckpointSaverHooks' Saver and filter out.
if isinstance(hook, tf_compat.CheckpointSaverHook):
hook._saver = None # pylint: disable=protected-access
continue
# Do not decorate the _OverwriteCheckpointHook since it should always
# run during the graph growing phase.
if not isinstance(hook, _OverwriteCheckpointHook):
hook = _GraphGrowingHookDecorator(hook)
processed_hooks.append(hook)
return processed_hooks
def _training_chief_hooks(self, current_iteration, training):
"""Returns chief-only training hooks to be run this iteration.
Args:
current_iteration: Current `_Iteration`.
training: Whether in training mode.
Returns:
A list of `SessionRunHook` instances.
"""
if not training:
return []
training_hooks = []
if tf_compat.is_v2_behavior_enabled():
# Use V2 summaries and hook when user is using TF 2 behavior.
training_hooks.append(
_SummaryV2SaverHook(
current_iteration.summaries,
save_steps=self.config.save_summary_steps))
else:
# Fallback to V1 summaries.
for summary in current_iteration.summaries:
output_dir = self.model_dir
if summary.scope:
output_dir = os.path.join(output_dir, summary.namespace,
summary.scope)
summary_saver_hook = tf_compat.SummarySaverHook(
save_steps=self.config.save_summary_steps,
output_dir=output_dir,
summary_op=summary.merge_all())
training_hooks.append(summary_saver_hook)
training_hooks += list(
current_iteration.estimator_spec.training_chief_hooks)
return training_hooks
def _training_hooks(self, current_iteration, training,
iteration_number_tensor, previous_iteration_vars,
is_growing_phase):
"""Returns training hooks to be run on all workers and chief this iteration.
Args:
current_iteration: Current `_Iteration`.
training: Whether in training mode.
iteration_number_tensor: An int tensor of the current AdaNet iteraiton.
previous_iteration_vars: The variables of the previous iteration to be
restored by the _OverwriteCheckpointHook. If empty, no
_OverwriteCheckpointHook will be created.
is_growing_phase: Whether we are in the AdaNet graph growing phase.
Returns:
A list of `SessionRunHook` instances.
"""
if not training:
return []
def after_fn():
self._iteration_ended = True
training_hooks = list(current_iteration.estimator_spec.training_hooks) + [
_StopAfterTrainingHook(current_iteration, after_fn=after_fn)
]
if is_growing_phase:
training_hooks.append(
_OverwriteCheckpointHook(current_iteration, iteration_number_tensor,
previous_iteration_vars, self.config,
self._enable_v2_checkpoint))
return training_hooks
def _evaluation_hooks(self, current_iteration, training, evaluation_name):
"""Returns evaluation hooks for this iteration.
Args:
current_iteration: Current `_Iteration`.
training: Whether in training mode.
evaluation_name: String name to append to the eval directory.
Returns:
A list of `SessionRunHook` instances.
"""
if training:
return []
evaluation_hooks = []
for subnetwork_spec in current_iteration.subnetwork_specs:
evaluation_hooks.append(
self._create_eval_metric_saver_hook(
subnetwork_spec.eval_metrics,
subnetwork_spec.name,
kind="subnetwork",
evaluation_name=evaluation_name))
for candidate in current_iteration.candidates:
evaluation_hooks.append(
self._create_eval_metric_saver_hook(
candidate.ensemble_spec.eval_metrics,
candidate.ensemble_spec.name,
kind="ensemble",
evaluation_name=evaluation_name))
return evaluation_hooks
def _create_eval_metric_saver_hook(self, eval_metrics, name, kind,
evaluation_name):
eval_subdir = "eval"
if evaluation_name:
eval_subdir = "eval_{}".format(evaluation_name)
return _EvalMetricSaverHook(
name=name,
kind=kind,
eval_metrics=eval_metrics,
output_dir=os.path.join(self.model_dir, kind, name, eval_subdir))
def _save_architecture(self, filename, architecture, checkpoint_path):
# type: (Text, _Architecture, Text) -> None
"""Persists the ensemble's architecture in a serialized format.
Writes to a text file with one subnetwork's iteration number and name
per line.
Args:
filename: String filename to persist the ensemble architecture.
architecture: Target `_Architecture` instance.
checkpoint_path: Path of the checkpoint to use.
"""
# Make directories since model_dir may not have been created yet.
tf.io.gfile.makedirs(os.path.dirname(filename))
iteration_number = self._checkpoint_iteration_number(checkpoint_path)
global_step = self._checkpoint_global_step(checkpoint_path)
serialized_architecture = architecture.serialize(iteration_number,
global_step)
logging.info("Saving architecture to %s:\n%s", filename,
serialized_architecture)
with tf.io.gfile.GFile(filename, "w") as record_file:
record_file.write(serialized_architecture)
def _read_architecture(self, filename):
# type: (Text) -> _Architecture
"""Reads an ensemble architecture from disk.
Assumes the file was written with `_save_architecture`.
Args:
filename: String filename where features were recorded.
Returns:
An `_Architecture` instance.
Raises:
OSError: When file not found at `filename`.
"""
if not tf.io.gfile.exists(filename):
raise OSError(errno.ENOENT, os.strerror(errno.ENOENT), filename)
with tf.io.gfile.GFile(filename, "rb") as gfile:
return _Architecture.deserialize(gfile.read().decode())
def _find_ensemble_candidate(self, ensemble_candidate_name,
ensemble_candidates):
# type: (Text, Sequence[ensemble_lib.Candidate]) -> ensemble_lib.Candidate
"""Returns the ensemble candidate with the given name."""
for ensemble_candidate in ensemble_candidates:
if ensemble_candidate.name == ensemble_candidate_name:
return ensemble_candidate
raise ValueError(
"Could not find a matching ensemble candidate with name '{}'. "
"Are you sure the `adanet.ensemble.Strategy` is deterministic?".format(
ensemble_candidate_name))
# TODO: Refactor architecture building logic to its own module.
def _architecture_ensemble_spec(self, architecture, iteration_number,
features, mode, labels,
previous_ensemble_spec, config,
previous_iteration, hooks):
"""Returns an `_EnsembleSpec` with the given architecture.
Creates the ensemble architecture by calling `generate_subnetworks` on
`self._subnetwork_generator` and only calling `build_subnetwork` on
`Builders` included in the architecture. Once their ops are created, their
variables are restored from the checkpoint.
Args:
architecture: An `_Architecture` instance.
iteration_number: Integer current iteration number.
features: Dictionary of `Tensor` objects keyed by feature name.
mode: Defines whether this is training, evaluation or prediction. See
`ModeKeys`.
labels: Labels `Tensor` or a dictionary of string label name to `Tensor`
(for multi-head). Can be `None`.
previous_ensemble_spec: The `_EnsembleSpec` for the previous iteration.
Will be `None` for the first iteration.
config: The current `tf.estimator.RunConfig`.
previous_iteration: The previous `_Iteration`.
hooks: A list of `tf.estimator.SessionRunHook`s.
Returns:
An `EnsembleSpec` instance for the given architecture.
Raises:
ValueError: If a subnetwork from `architecture` is not found in the
generated candidate `Builders` of the specified iteration.
"""
previous_ensemble = None
if previous_ensemble_spec:
previous_ensemble = previous_ensemble_spec.ensemble
current_iteration = previous_iteration
for t, names in architecture.subnetworks_grouped_by_iteration:
if t != iteration_number:
continue
previous_ensemble_reports, all_reports = [], []
if self._report_materializer:
previous_ensemble_reports, all_reports = (
self._collate_subnetwork_reports(iteration_number))
generated_subnetwork_builders = (
self._call_generate_candidates(
previous_ensemble=previous_ensemble,
iteration_number=iteration_number,
previous_ensemble_reports=previous_ensemble_reports,
all_reports=all_reports,
config=config))
subnetwork_builder_names = {
b.name: b for b in generated_subnetwork_builders
}
rebuild_subnetwork_builders = []
for name in names:
if name not in subnetwork_builder_names:
raise ValueError(
"Required subnetwork builder is missing for iteration {}: {}"
.format(iteration_number, name))
rebuild_subnetwork_builders.append(subnetwork_builder_names[name])
previous_ensemble_summary = None
previous_ensemble_subnetwork_builders = None
if previous_ensemble_spec:
# Always skip summaries when rebuilding previous architecture,
# since they are not useful.
previous_ensemble_summary = self._summary_maker(
namespace="ensemble",
scope=previous_ensemble_spec.name,
skip_summary=True)
previous_ensemble_subnetwork_builders = (
previous_ensemble_spec.subnetwork_builders)
ensemble_candidates = []
for ensemble_strategy in self._ensemble_strategies:
ensemble_candidates += ensemble_strategy.generate_ensemble_candidates(
rebuild_subnetwork_builders, previous_ensemble_subnetwork_builders)
ensemble_candidate = self._find_ensemble_candidate(
architecture.ensemble_candidate_name, ensemble_candidates)
current_iteration = self._iteration_builder.build_iteration(
base_global_step=architecture.global_step,
iteration_number=iteration_number,
ensemble_candidates=[ensemble_candidate],
subnetwork_builders=rebuild_subnetwork_builders,
features=features,
labels=labels,
mode=mode,
config=config,
previous_ensemble_summary=previous_ensemble_summary,
rebuilding=True,
rebuilding_ensembler_name=architecture.ensembler_name,
previous_iteration=current_iteration)
max_candidates = 2 if previous_ensemble_spec else 1
assert len(current_iteration.candidates) == max_candidates
previous_ensemble_spec = current_iteration.candidates[-1].ensemble_spec
previous_ensemble = previous_ensemble_spec.ensemble
previous_ensemble_spec.architecture.set_replay_indices(
architecture.replay_indices)
return current_iteration
def _collate_subnetwork_reports(self, iteration_number):
"""Prepares subnetwork.Reports to be passed to Generator.
Reads subnetwork.MaterializedReports from past iterations,
collates those that were included in previous_ensemble into
previous_ensemble_reports as a List of subnetwork.MaterializedReports,
and collates all reports from previous iterations into all_reports as
another List of subnetwork.MaterializedReports.
Args:
iteration_number: Python integer AdaNet iteration number, starting from 0.
Returns:
(previous_ensemble_reports: List<subnetwork.MaterializedReport>,
materialized_reports: List<MaterializedReport>)
"""
materialized_reports_all = (self._report_accessor.read_iteration_reports())
previous_ensemble_reports = []
all_reports = []
# Since the number of iteration reports changes after the
# MATERIALIZE_REPORT phase, we need to make sure that we always pass the
# same reports to the Generator in the same iteration,
# otherwise the graph that is built in the FREEZE_ENSEMBLE phase would be
# different from the graph built in the training phase.
# Iteration 0 should have 0 iteration reports passed to the
# Generator, since there are no previous iterations.
# Iteration 1 should have 1 list of reports for Builders
# generated in iteration 0.
# Iteration 2 should have 2 lists of reports -- one for iteration 0,
# one for iteration 1. Note that the list of reports for iteration >= 1
# should contain "previous_ensemble", in addition to the
# Builders at the start of that iteration.
# Iteration t should have t lists of reports.
for i, iteration_reports in enumerate(materialized_reports_all):
# This ensures that the FREEZE_ENSEMBLE phase does not pass the reports
# generated in the previous phase of the same iteration to the
# Generator when building the graph.
if i >= iteration_number:
break
chosen_subnetworks_in_this_iteration = [
subnetwork_report for subnetwork_report in iteration_reports
if subnetwork_report.included_in_final_ensemble
]
previous_ensemble_reports += chosen_subnetworks_in_this_iteration
all_reports.extend(iteration_reports)
return previous_ensemble_reports, all_reports
def _train_op(self, iteration_estimator_spec, is_growing_phase):
"""Returns the iteration train op or tf.no_op if growing the graph."""
train_op = iteration_estimator_spec.train_op
if is_growing_phase:
train_op = tf_compat.v1.train.get_global_step().assign_add(1)
# NOTE: some version of TensorFlow check that train_op is an Op or Tensor
# and crash if train_op is a Variable.
train_op = tf.identity(train_op)
return train_op
def _create_estimator_spec(self, current_iteration, mode,
iteration_number_tensor, previous_iteration_vars,
is_growing_phase, evaluation_name):
"""Creates the EstimatorSpec which will be returned by _adanet_model_fn."""
from tensorflow.python.training.tracking import graph_view # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
training = mode == tf.estimator.ModeKeys.TRAIN
iteration_estimator_spec = current_iteration.estimator_spec
training_chief_hooks = self._training_chief_hooks(current_iteration,
training)
training_hooks = self._training_hooks(current_iteration, training,
iteration_number_tensor,
previous_iteration_vars,
is_growing_phase)
if is_growing_phase:
training_chief_hooks = self._process_hooks_for_growing_phase(
training_chief_hooks)
training_hooks = self._process_hooks_for_growing_phase(training_hooks)
saver = None
if self._enable_v2_checkpoint:
saver = tf_compat.v1.train.Saver(
var_list=graph_view.ObjectGraphView(
current_iteration.checkpoint).frozen_saveable_objects(),
sharded=True,
max_to_keep=self.config.keep_checkpoint_max)
return tf.estimator.EstimatorSpec(
mode=mode,
predictions=iteration_estimator_spec.predictions,
loss=iteration_estimator_spec.loss,
train_op=self._train_op(iteration_estimator_spec, is_growing_phase),
eval_metric_ops=iteration_estimator_spec.eval_metric_ops,
training_chief_hooks=training_chief_hooks,
training_hooks=training_hooks,
evaluation_hooks=self._evaluation_hooks(current_iteration, training,
evaluation_name),
scaffold=tf_compat.v1.train.Scaffold(
summary_op=tf.constant(""),
saver=saver,
local_init_op=current_iteration.estimator_spec.scaffold
.local_init_op if isinstance(current_iteration.estimator_spec,
tf.estimator.EstimatorSpec) else None),
export_outputs=iteration_estimator_spec.export_outputs)
def _call_generate_candidates(self, previous_ensemble, iteration_number,
previous_ensemble_reports, all_reports, config):
# Calling low level getargs for py_2_and_3 compatibility.
defined_args = inspect.getargs(
self._subnetwork_generator.generate_candidates.__code__).args
generate_args = dict(
previous_ensemble=previous_ensemble,
iteration_number=iteration_number,
previous_ensemble_reports=previous_ensemble_reports,
all_reports=all_reports)
if "config" in defined_args:
generate_args["config"] = config
return self._subnetwork_generator.generate_candidates(**generate_args)
def _create_iteration(self,
features,
labels,
mode,
config,
is_growing_phase,
checkpoint_path,
hooks,
best_ensemble_index_override=None):
"""Constructs the TF ops and variables for the current iteration.
Args:
features: Dictionary of `Tensor` objects keyed by feature name.
labels: Labels `Tensor` or a dictionary of string label name to `Tensor`
(for multi-head). Can be `None`.
mode: Defines whether this is training, evaluation or prediction. See
`ModeKeys`.
config: The current `tf.estimator.RunConfig`.
is_growing_phase: Whether we are in the AdaNet graph growing phase.
checkpoint_path: Path of the checkpoint to use. When `None`, this method
uses the latest checkpoint instead.
hooks: A list of `tf.estimator.SessionRunHooks`.
best_ensemble_index_override: Integer index to identify the latest
iteration's best ensemble candidate instead of computing the best
ensemble index dynamically conditional on the ensemble AdaNet losses.
Returns:
A two-tuple of the current `_Iteration`, and list of variables from
the previous iteration for restoring during the graph growing phase.
"""
# Use the evaluation checkpoint path to get both the iteration number and
# variable values to avoid any race conditions between the first and second
# checkpoint reads.
iteration_number = self._checkpoint_iteration_number(checkpoint_path)
if mode == tf.estimator.ModeKeys.EVAL and checkpoint_path is None:
# This should only happen during some tests, so we log instead of
# asserting here.
logging.warning("There are no checkpoints available during evaluation. "
"Variables will be initialized to their defaults.")
if is_growing_phase:
assert mode == tf.estimator.ModeKeys.TRAIN
assert config.is_chief
iteration_number += 1
# Only record summaries when training.
skip_summaries = (mode != tf.estimator.ModeKeys.TRAIN or is_growing_phase)
base_global_step = 0
with tf_compat.v1.variable_scope("adanet"):
previous_iteration = None
previous_ensemble_spec = None
previous_ensemble = None
previous_ensemble_summary = None
previous_ensemble_subnetwork_builders = None
architecture = None
for i in range(iteration_number):
architecture_filename = self._architecture_filename(i)
if not tf.io.gfile.exists(architecture_filename):
continue
architecture = self._read_architecture(architecture_filename)
logging.info(
"Importing architecture from %s: [%s].", architecture_filename,
", ".join(
sorted([
"'{}:{}'".format(t, n)
for t, n in architecture.subnetworks_grouped_by_iteration
])))
base_global_step = architecture.global_step
previous_iteration = self._architecture_ensemble_spec(
architecture, i, features, mode, labels, previous_ensemble_spec,
config, previous_iteration, hooks)
previous_ensemble_spec = previous_iteration.candidates[-1].ensemble_spec
previous_ensemble = previous_ensemble_spec.ensemble
previous_ensemble_summary = self._summary_maker(
namespace="ensemble",
scope=previous_ensemble_spec.name,
skip_summary=skip_summaries)
previous_ensemble_subnetwork_builders = (
previous_ensemble_spec.subnetwork_builders)
previous_iteration_vars = None
if is_growing_phase:
# Keep track of the previous iteration variables so we can restore them
# from the previous checkpoint after growing the graph. After this line,
# any variables created will not have a matching one in the checkpoint
# until it gets overwritten.
# Note: It's not possible to just create a tf.train.Saver here since
# this code is also run on TPU, which does not support creating Savers
# inside model_fn.
previous_iteration_vars = (
tf_compat.v1.get_collection(tf_compat.v1.GraphKeys.GLOBAL_VARIABLES)
+ tf_compat.v1.get_collection(
tf_compat.v1.GraphKeys.SAVEABLE_OBJECTS))
previous_ensemble_reports, all_reports = [], []
if self._report_materializer:
previous_ensemble_reports, all_reports = (
self._collate_subnetwork_reports(iteration_number))
subnetwork_builders = self._call_generate_candidates(
previous_ensemble=previous_ensemble,
iteration_number=iteration_number,
previous_ensemble_reports=previous_ensemble_reports,
all_reports=all_reports,
config=config)
ensemble_candidates = []
for ensemble_strategy in self._ensemble_strategies:
ensemble_candidates += ensemble_strategy.generate_ensemble_candidates(
subnetwork_builders, previous_ensemble_subnetwork_builders)
current_iteration = self._iteration_builder.build_iteration(
base_global_step=base_global_step,
iteration_number=iteration_number,
ensemble_candidates=ensemble_candidates,
subnetwork_builders=subnetwork_builders,
features=features,
labels=labels,
mode=mode,
config=config,
previous_ensemble_summary=previous_ensemble_summary,
best_ensemble_index_override=best_ensemble_index_override,
previous_iteration=previous_iteration)
return current_iteration, previous_iteration_vars
def _create_model_fn(self,
is_growing_phase=False,
is_inside_training_loop=False,
is_export=False,
evaluation_name=None,
best_ensemble_index=None,
checkpoint_path=None,
hooks=None):
"""Creates the AdaNet model_fn.
Args:
is_growing_phase: Whether the model_fn will be called in the growing
phase.
is_inside_training_loop: Whether the model_fn will be called inside the
AdaNet training loop.
is_export: Whether the model_fn will be called from functions which export
a SavedModel.
evaluation_name: String name to append to the eval directory.
best_ensemble_index: The index of the best performing ensemble in the
latest AdaNet iteration.
checkpoint_path: The checkpoint path from which to restore variables.
hooks: Extra hooks to use when creating the graph.
Returns:
The adanet_model_fn which will create the computation graph when called.
"""
del is_export # Unused.
def _adanet_model_fn(features, labels, mode, params, config):
"""AdaNet model_fn.
Args:
features: Dictionary of `Tensor` objects keyed by feature name.
labels: Labels `Tensor` or a dictionary of string label name to `Tensor`
(for multi-head). Can be `None`.
mode: Defines whether this is training, evaluation or prediction. See
`ModeKeys`.
params: A dict of parameters.
config: The current `tf.estimator.RunConfig`.
Returns:
A `EstimatorSpec` instance.
Raises:
UserWarning: When calling model_fn directly in TRAIN mode.
"""
del params # Unused.
path = checkpoint_path or tf.train.latest_checkpoint(self.model_dir)
training = mode == tf.estimator.ModeKeys.TRAIN
if training and not is_inside_training_loop:
raise UserWarning(
"The adanet.Estimator's model_fn should not be called directly in "
"TRAIN mode, because its behavior is undefined outside the context "
"of its `train` method. If you are trying to add custom metrics "
"with `tf.contrib.estimator.add_metrics`, pass the `metric_fn` to "
"this `Estimator's` constructor instead.")
current_iteration, previous_iteration_vars = self._create_iteration(
features,
labels,
mode,
config,
is_growing_phase,
checkpoint_path=path,
hooks=hooks,
best_ensemble_index_override=best_ensemble_index)
# Variable which allows us to read the current iteration from a
# checkpoint. This must be created here so it is available when calling
# _execute_bookkeeping_phase after the first iteration.
iteration_number_tensor = None
if not self._enable_v2_checkpoint:
iteration_number_tensor = tf_compat.v1.get_variable(
self._Keys.CURRENT_ITERATION,
shape=[],
dtype=tf.int64,
initializer=tf_compat.v1.zeros_initializer(),
trainable=False)
return self._create_estimator_spec(
current_iteration,
mode,
iteration_number_tensor,
previous_iteration_vars,
is_growing_phase,
evaluation_name=evaluation_name)
return _adanet_model_fn