controller.py 22.6 KB
Newer Older
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1
# Copyright 2022 The Orbit Authors. All Rights Reserved.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
2
3
4
5
6
7
8
9
10
11
12
13
#
# 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
#
#     http://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.
Hongkun Yu's avatar
Hongkun Yu committed
14

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
15
"""Provides a `Controller` class for managing the outer training loop."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
16

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
17
import pprint
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
18
import time
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
19

Jiayu Ye's avatar
Jiayu Ye committed
20
from typing import Any, Callable, Iterable, Optional, Union
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
21

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
22
from absl import logging
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
23

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
24
25
26
27
28
29
from orbit import runner
from orbit import utils

import tensorflow as tf


Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
30
def _log(message: str):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
31
32
33
34
35
  """Logs `message` to the `info` log, and also prints to stdout."""
  logging.info(message)
  print(message)


Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
36
37
38
39
40
41
42
43
44
45
46
47
48
logging.ABSLLogger.register_frame_to_skip(__file__, _log.__name__)


def _format_output(output, indent=4):
  """Formats `output`, either on one line, or indented across multiple lines."""
  formatted = pprint.pformat(output)
  lines = formatted.splitlines()
  if len(lines) == 1:
    return formatted
  lines = [" " * indent + line for line in lines]
  return "\n" + "\n".join(lines)


Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
49
50
51
Action = Callable[[runner.Output], None]


Hongkun Yu's avatar
Hongkun Yu committed
52
class Controller:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
53
54
55
56
57
58
  """Class that controls the outer loop of model training and evaluation.

  Orbit divides training and evaluation into "inner" and "outer" loops. Inner
  loops are implemented by users in the form of `AbstractTrainer` and
  `AbstractEvaluator` subclasses, and define how to run a given number of
  training or evaluation steps. The outer loop is provided by this `Controller`,
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
59
60
61
  and interleaves calls to the user-provided inner loops with additional actions
  such as saving checkpoints, running evaluations, writing summaries, as well as
  (optionally) user provided `Action`s (see below).
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
62
63
64
65
66
67
68
69
70
71
72
73
74

  There are four top-level "outer loops" provided:

    - `train`, which trains until a specified number of global steps is reached;
    - `evaluate`, for one-off model evaluation;
    - `train_and_evaluate`, for interleaved training and evaluation;
    - `evaluate_continuously`, for monitoring a given directory and running
      evaluations on new model checkpoints.

  While this class attempts to provide out-of-the-box solutions for common
  training and evaluation use cases, the internal details and method
  implementations are also intended to be simple enough to make subclassing or
  other custom outer loop implementations easy to achieve.
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
75
76

  Some additional customization can be achieved by supplying `train_actions` or
Ron Shapiro's avatar
Ron Shapiro committed
77
78
79
80
81
82
83
  `eval_actions` when constructing the `Controller`. Actions arbitrary callables
  that are applied by the `Controller` to the output of train steps (after each
  inner loop of `steps_per_loop` steps) or an evaluation. This provides a hook
  mechanism, enabling things like reporting metrics to Vizier, model exporting,
  additional logging, etc. See the `orbit.actions` package for a small handful
  of predefined actions and some utility classes that may be useful in defining
  your own.
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
84
  """
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
85
86
87

  def __init__(
      self,
88
89
      *,  # Makes all args keyword only.
      global_step: tf.Variable,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
90
91
      trainer: Optional[runner.AbstractTrainer] = None,
      evaluator: Optional[runner.AbstractEvaluator] = None,
92
      strategy: Optional[tf.distribute.Strategy] = None,
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
93
      # Actions
Ron Shapiro's avatar
Ron Shapiro committed
94
95
      train_actions: Optional[Iterable[Action]] = None,
      eval_actions: Optional[Iterable[Action]] = None,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
96
      # Train related
97
      steps_per_loop: Optional[Union[int, Callable[[int], int]]] = None,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
98
99
100
      checkpoint_manager: Optional[tf.train.CheckpointManager] = None,
      # Summary related
      summary_interval: Optional[int] = None,
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
101
      summary_dir: Optional[str] = None,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
102
      # Evaluation related
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
103
      eval_summary_dir: Optional[str] = None,
Jiayu Ye's avatar
Jiayu Ye committed
104
105
      summary_manager: Optional[Any] = None,
      eval_summary_manager: Optional[Any] = None):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
106
107
108
109
110
    """Initializes a `Controller` instance.

    Note that if `checkpoint_manager` is provided and there are checkpoints in
    the associated model directory, the model will be restored from the most
    recent checkpoint during this `__init__` method.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
111
112

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
113
114
115
116
117
118
119
120
121
      global_step: An integer `tf.Variable` storing the global training step
        number. Usually this can be obtained from the `iterations` property of
        the model's optimizer (e.g. `trainer.optimizer.iterations`). In cases
        where multiple optimizers are used, or if one model "step" corresponds
        to more than one update to model parameters, users can create and
        increment their own global step variable as well. In this case it is
        recommended to create the `tf.Variable` inside the distribution strategy
        scope, with `aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA` (see
        also `orbit.utils.create_global_step()`).
122
123
124
125
126
127
128
      trainer: An instance of `orbit.AbstractTrainer`, which implements the
        inner training loop.
      evaluator: An instance of `orbit.AbstractEvaluator`, which implements
        evaluation.
      strategy: An instance of `tf.distribute.Strategy`. If not provided, the
        strategy will be initialized from the current in-scope strategy using
        `tf.distribute.get_strategy()`.
Ron Shapiro's avatar
Ron Shapiro committed
129
130
131
132
133
      train_actions: Optional `orbit.Action`s to call after each block of
        `steps_per_loop` training steps are run. These will be called with the
        output of `trainer.train`.
      eval_actions: Optional `orbit.Action`s to call after each evaluation.
        These will be called with the output of `evaluator.evaluate`.
134
135
136
137
138
      steps_per_loop: Optional integer to indicate the number of steps to run in
        each inner loop of training (passed as the `num_steps` parameter of
        `trainer.train`). It can be also a callable which takes the current
        global step value as input and returns the number of steps to run as
        output.
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
139
140
141
142
143
      checkpoint_manager: An instance of `tf.train.CheckpointManager`. If
        provided and there are checkpoints in the associated model directory,
        the model will be restored from the most recent checkpoint inside this
        `__init__` method. If not provided, the `Controller` will not
        automatically save to or restore from checkpoints.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
144
      summary_interval: Step interval for training summaries. Note that this
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
145
146
147
148
149
150
151
152
153
154
155
        argument only applies to `tf.summary` calls inside the `trainer.train`
        function. Summaries written by the `Controller` (specifically
        "steps_per_second" and output from the `trainer.train` method) will
        always be enabled unless the `summary_dir` parameter is `None`. If set,
        the value must be divisible by `steps_per_loop`.
      summary_dir: The directory to write summaries to. To use the same
        directory as for checkpointing, pass `checkpoint_manager.directory`. If
        `None`, no training summaries will be written.
      eval_summary_dir: The directory to write eval summaries to. If `None`, it
        will be set to `summary_dir`. If both `summary_dir` and
        `eval_summary_dir` are `None`, no eval summaries will be written.
Jiayu Ye's avatar
Jiayu Ye committed
156
157
158
159
160
161
162
163
      summary_manager: Instance of the summary manager. If set, the
        `summary_dir` will be ignored. Otherwise the summary manager will be
        created internally for TensorBoard summaries by default from the
        `summary_dir`.
      eval_summary_manager: Instance of the eval summary manager. If set, the
        `eval_summary_dir` will be ignored. Otherwise the eval summary manager
        will be created internally for TensorBoard summaries by default from the
        `eval_summary_dir`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
164
165

    Raises:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
166
      ValueError: If both `trainer` and `evaluator` are `None`.
167
      ValueError: If `steps_per_loop` is not a positive integer or a callable.
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
168
169
      ValueError: If `summary_interval` is not a positive integer or is not
        divisible by `steps_per_loop`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
170
171
    """
    if trainer is None and evaluator is None:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
172
      raise ValueError("`trainer` and `evaluator` should not both be `None`.")
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
173

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
174
175
    if trainer is not None:
      if steps_per_loop is None:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
176
177
        raise ValueError(
            "`steps_per_loop` is required when `trainer` is provided.")
178
179
      elif not callable(steps_per_loop) and (
          not isinstance(steps_per_loop, int) or steps_per_loop < 1):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
180
        raise ValueError(
181
182
            f"`steps_per_loop` ({steps_per_loop}) must be a positive integer "
            "or a callable.")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
183
184
185

      if summary_interval is not None:
        if summary_interval <= 0:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
186
187
          raise ValueError(
              f"`summary_interval` ({summary_interval}) must be larger than 0.")
188
189
        elif not callable(steps_per_loop) and (summary_interval % steps_per_loop
                                               != 0):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
190
191
192
193
          raise ValueError(
              f"`summary interval` ({summary_interval}) must be a multiple "
              f"of `steps_per_loop` ({steps_per_loop}).")

194
    if not isinstance(global_step, tf.Variable):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
195
      raise ValueError("`global_step` must be a `tf.Variable`.")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
196
197
198
199
200
201

    self.trainer = trainer
    self.evaluator = evaluator

    self.strategy = strategy or tf.distribute.get_strategy()

Ron Shapiro's avatar
Ron Shapiro committed
202
203
    self.train_actions = () if train_actions is None else tuple(train_actions)
    self.eval_actions = () if eval_actions is None else tuple(eval_actions)
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
204

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
205
206
207
208
209
210
    self.global_step = global_step
    self.checkpoint_manager = checkpoint_manager

    if self.trainer is not None:
      self.step_timer = None
      self.summary_interval = summary_interval
Jiayu Ye's avatar
Jiayu Ye committed
211
212
213
214
215
      if summary_manager:
        self.summary_manager = summary_manager
      else:
        self.summary_manager = utils.SummaryManager(
            summary_dir, tf.summary.scalar, global_step=self.global_step)
216
      self._steps_per_loop = steps_per_loop
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
217
218
219
220
221
222
223
224

    if self.evaluator is not None:
      eval_summary_dir = eval_summary_dir or summary_dir
      if eval_summary_dir == summary_dir and self.trainer is not None:
        # Reuse the summary writer if train and evaluation summary directory
        # are the same.
        self.eval_summary_manager = self.summary_manager
      else:
Jiayu Ye's avatar
Jiayu Ye committed
225
226
227
228
229
        if eval_summary_manager:
          self.eval_summary_manager = eval_summary_manager
        else:
          self.eval_summary_manager = utils.SummaryManager(
              eval_summary_dir, tf.summary.scalar, global_step=self.global_step)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
230

231
    tf.summary.experimental.set_step(self.global_step)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
232
233
234

    # Restores the model if needed.
    if self.checkpoint_manager is not None:
235
236
      restored_path = self.restore_checkpoint()
      if restored_path:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
237
        _log(f"restored from checkpoint: {restored_path}")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
238
239

  def train(self, steps: int, checkpoint_at_completion: bool = True):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
240
    """Runs training until the specified global step count has been reached.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
241

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
242
243
244
245
    This method makes calls to `self.trainer.train()` until the global step
    count is equal to `steps`. It will additionally save checkpoints (if a
    `CheckpointManager` was passed to `Controller.__init__`) and summarize
    training output (if `summary_dir` is set).
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
246
247
248
249

    Args:
      steps: The global step count to train up to.
      checkpoint_at_completion: Whether to save a checkpoint when this method
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
250
        returns (regardless of the checkpointing interval). Defaults to `True`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
251
    """
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
252
    self._require("trainer", for_method="train")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
253
254

    # TODO(momernick): Support steps=None or -1 (training to exhaustion).
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
255
256
    current_step = self.global_step.numpy()  # Cache, since this is expensive.
    _log(f"train | step: {current_step: 6d} | training until step {steps}...")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
257
258
259
260
261
    while current_step < steps:
      # Calculates steps to run for the next train loop.
      num_steps = min(steps - current_step, self.steps_per_loop)
      self._train_n_steps(num_steps)
      self._maybe_save_checkpoint()
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
262
      current_step = self.global_step.numpy()
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
263
264

    if checkpoint_at_completion:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
265
      self._maybe_save_checkpoint(check_interval=False)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
266

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
267
268
  def evaluate(self, steps: int = -1) -> Optional[runner.Output]:
    """Runs evaluation for the given number of steps.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
269

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
270
271
    This method calls `self.evaluator.evaluate(steps)`, then writes the returned
    summaries (if any).
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
272
273

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
274
275
276
277
      steps: The number of evaluation steps to run. The value `-1` is reserved
        as a special sentinel to indicate a "complete" evaluation that runs
        until the underlying dataset is exhausted. Support for this is dependent
        on the specific `evaluator` being used.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
278

Simon Kornblith's avatar
Simon Kornblith committed
279
    Returns:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
280
      The evaluation results as a dictionary mapping names to NumPy values.
Simon Kornblith's avatar
Simon Kornblith committed
281

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
282
    Raises:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
283
284
285
      ValueError: If `evaluator` was not provided to `Controller.__init__`.
      ValueError: If no checkpoint is present in `checkpoint_manager.directory`.
      ValueError: If `steps` is not a positive value or -1.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
286
    """
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
287
    self._require("evaluator", for_method="evaluate")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
288
289

    if steps > 0:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
290
291
292
      steps_msg = f"running {steps} steps of evaluation..."
    elif steps == -1:
      steps_msg = "running complete evaluation..."
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
293
    else:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
294
      raise ValueError(f"`steps` ({steps}) should be > 0, or == -1.")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
295

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
296
297
    current_step = self.global_step.numpy()
    _log(f" eval | step: {current_step: 6d} | {steps_msg}")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
298

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
299
300
301
302
303
    start = time.time()
    with self.eval_summary_manager.summary_writer().as_default():
      steps_tensor = tf.convert_to_tensor(steps, dtype=tf.int32)
      eval_output = self.evaluator.evaluate(steps_tensor)
    elapsed = time.time() - start
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
304

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
305
306
307
308
309
    eval_output = eval_output or {}
    for action in self.eval_actions:
      action(eval_output)
    eval_output = tf.nest.map_structure(utils.get_value, eval_output)

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
310
    _log(f" eval | step: {current_step: 6d} | "
311
         f"eval time: {elapsed: 6.1f} sec | "
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
312
         f"output: {_format_output(eval_output)}")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
313

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
314
    self.eval_summary_manager.write_summaries(eval_output)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
315
316
    self.eval_summary_manager.flush()

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
317
    return eval_output
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
318
319

  def train_and_evaluate(self,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
320
                         train_steps: int,
321
                         eval_steps: int = -1,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
322
                         eval_interval: Optional[int] = None) -> None:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
323
    """Runs interleaved training and evaluation.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
324

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
325
326
327
328
329
    This method interleaves calls to `self.train()` and `self.evaluate()`,
    training the model until the global step count equals `train_steps`, and
    running an evaluation for `eval_steps` every `eval_interval` training steps.
    In addition, this method will run a final evaluation at the end of the
    training sequence.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
330
331
332

    Args:
      train_steps: The global step count to train up to.
333
      eval_steps: The number of steps to run during an evaluation. If -1, this
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
334
335
336
337
        method will evaluate over the entire evaluation dataset.
      eval_interval: The number of training steps to run between evaluations. If
        set, training will always stop every `eval_interval` steps, even if this
        results in a shorter inner loop than specified by `steps_per_loop`
Ruoxin Sang's avatar
Ruoxin Sang committed
338
339
        setting. If None, evaluation will only be performed after training is
        complete.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
340
    """
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
341
342
343
344
    self._require("trainer", for_method="train_and_evaluate")
    self._require("evaluator", for_method="train_and_evaluate")

    current_step = self.global_step.numpy()  # Cache, since this is expensive.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
345
346
347
348
349
350
    eval_interval = eval_interval or (train_steps - current_step)
    while current_step < train_steps:
      interval = min(train_steps - current_step, eval_interval)
      num_steps = current_step + interval
      self.train(steps=num_steps, checkpoint_at_completion=False)
      self.evaluate(steps=eval_steps)
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
351
352
      current_step = self.global_step.numpy()
    self._maybe_save_checkpoint(check_interval=False)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
353
354

  def evaluate_continuously(self,
355
                            steps: int = -1,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
356
357
                            timeout: Optional[Union[int, float]] = None,
                            timeout_fn: Optional[Callable[[], bool]] = None):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
358
    """Continuously monitors a directory and evaluates new checkpoints in it.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
359
360
361
362
363
364

    This method continuously monitors a directory as specified by this
    Controller's CheckpointManager init arg and runs evaluation on the
    checkpoints found there.

    Args:
365
366
      steps: The number of steps to run when evaluating. If -1, this method will
        evaluate over the entire evaluation dataset.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
367
368
369
370
371
372
373
374
375
376
      timeout: The maximum number of seconds to wait between checkpoints. See
        tf.train.checkpoints_iterator documentation.
      timeout_fn: Optional callable to call after a timeout. If the function
        returns True, then it means that no new checkpoints will be generated
        and the iterator will exit.

    Raises:
      ValueError: If no checkpoint found in `self.checkpoint_manager.directory`.
      ValueError: If `evaluator` was not provided as a controller init arg.
    """
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
377
378
379
    self._require("evaluator", for_method="evaluate_continuously")
    self._require("checkpoint_manager", for_method="evaluate_continuously")

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
380
381
382
383
384
385
386
    for checkpoint_path in tf.train.checkpoints_iterator(
        self.checkpoint_manager.directory,
        timeout=timeout,
        timeout_fn=timeout_fn):
      self.restore_checkpoint(checkpoint_path)
      self.evaluate(steps)

Rebecca Chen's avatar
Rebecca Chen committed
387
  def restore_checkpoint(self, checkpoint_path: Optional[str] = None):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
    """Restores the model from a checkpoint.

    Args:
      checkpoint_path: An optional string specifying the checkpoint path to
        restore from. If `None`, will restore from the most recent checkpoint
        (or initialize the model using a custom `init_fn` if no checkpoints can
        be found) using `self.checkpoint_manager.restore_or_initialize()`.

    Returns:
      The path to the restored checkpoint if a restore happened, or `None` if no
      restore occurred.
    """
    self._require("checkpoint_manager", for_method="restore_checkpoint")

    with self.strategy.scope():
      # Checkpoint restoring should be inside scope (b/139450638).
      if checkpoint_path is not None:
        _log(f"restoring model from {checkpoint_path}...")
        self.checkpoint_manager.checkpoint.restore(checkpoint_path)
      else:
        _log("restoring or initializing model...")
        checkpoint_path = self.checkpoint_manager.restore_or_initialize()

    if checkpoint_path is not None:
      _log(f"restored model from {checkpoint_path}.")
    else:
      _log("initialized model.")

    return checkpoint_path

  def save_checkpoint(self):
    """Saves the model to a checkpoint.

    This method will save a checkpoint containing the current state of the
    model.

    Raises:
      ValueError: If no `checkpoint_manager` was provided to
        `Controller.__init__`.
    """
    self._require("checkpoint_manager", for_method="save_checkpoint")
    self._maybe_save_checkpoint(check_interval=False)

431
432
433
434
435
436
437
  @property
  def steps_per_loop(self):
    """Returns current steps_per_loop value in a training loop."""
    if callable(self._steps_per_loop):
      return self._steps_per_loop(self.global_step.numpy())
    return self._steps_per_loop

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
438
  def _train_n_steps(self, num_steps: int):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
439
    """Runs training for `num_steps` steps.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
440

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
441
442
443
    Also prints/logs updates about training progress, and summarizes training
    output (if output is returned from `self.trainer.train()`, and if
    `self.summary_dir` is set).
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
444
445

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
446
      num_steps: An integer specifying how many steps of training to run.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
447
448

    Raises:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
449
450
      RuntimeError: If `global_step` is not properly incremented by `num_steps`
        after calling `self.trainer.train(num_steps)`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
451
452
453
454
455
    """
    if not self.step_timer:
      self.step_timer = StepTimer(self.global_step)
    current_step = self.global_step.numpy()

Ruoxin Sang's avatar
Ruoxin Sang committed
456
    with self.summary_manager.summary_writer().as_default():
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
457
458
      should_record = False  # Allows static optimization in no-summary cases.
      if self.summary_interval:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
459
        # Create a predicate to determine when summaries should be written.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
460
461
        should_record = lambda: (self.global_step % self.summary_interval == 0)
      with tf.summary.record_if(should_record):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
462
463
464
465
466
467
        num_steps_tensor = tf.convert_to_tensor(num_steps, dtype=tf.int32)
        train_output = self.trainer.train(num_steps_tensor)

    # Verify that global_step was updated properly, then update current_step.
    expected_step = current_step + num_steps
    if self.global_step.numpy() != expected_step:
468
      message = (
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
469
470
471
          f"`trainer.train({num_steps})` did not update `global_step` by "
          f"{num_steps}. Old value was {current_step}, expected updated value "
          f"to be {expected_step}, but it was {self.global_step.numpy()}.")
472
      logging.warning(message)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
473

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
474
475
476
477
478
    train_output = train_output or {}
    for action in self.train_actions:
      action(train_output)
    train_output = tf.nest.map_structure(utils.get_value, train_output)

479
    current_step = self.global_step.numpy()
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
480
    steps_per_second = self.step_timer.steps_per_second()
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
481
482
483
484
485
486
487
    _log(f"train | step: {current_step: 6d} | "
         f"steps/sec: {steps_per_second: 6.1f} | "
         f"output: {_format_output(train_output)}")

    train_output["steps_per_second"] = steps_per_second
    self.summary_manager.write_summaries(train_output)
    self.summary_manager.flush()
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
488

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
489
490
  def _maybe_save_checkpoint(self, check_interval: bool = True):
    """Conditionally saves a checkpoint.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
491

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
492
493
494
    A checkpoint is saved if a `CheckpointManager` is available, and if the
    required number of steps has elapsed since the last checkpoint was saved
    (although this condition can be disabled by setting `check_interval=False`).
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
495
496

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
497
498
499
500
      check_interval: Whether to check if the checkpoint interval has fully
        elapsed. If `False`, a checkpoint is saved regardless of the elapsed
        steps since the most recent checkpoint, unless no `checkpoint_manager`
        was provided to `Controller.__init__`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
501
502
503
504
505
506
507

    Returns:
      A boolean indicating whether a checkpoint was saved.
    """
    if self.checkpoint_manager and self.checkpoint_manager.checkpoint_interval:
      ckpt_path = self.checkpoint_manager.save(
          checkpoint_number=self.global_step.numpy(),
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
508
          check_interval=check_interval)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
509
      if ckpt_path is not None:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
510
        _log(f"saved checkpoint to {ckpt_path}.")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
511
512
513
        return True
    return False

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
514
515
516
517
518
519
520
  def _require(self, attribute, for_method):
    """Utility method to raise an error if the given `attribute` is not set."""
    if getattr(self, attribute, None) is None:
      raise ValueError(
          f"`{attribute}` is not set. Pass `{attribute}` to "
          f"`Controller.__init__` before calling `{for_method}()`.")

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
521

Hongkun Yu's avatar
Hongkun Yu committed
522
class StepTimer:
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
  """Utility class for measuring steps/second."""

  def __init__(self, step):
    self.step = step
    self.start()

  def start(self):
    self.last_iteration = self.step.numpy()
    self.last_time = time.time()

  def steps_per_second(self, restart=True):
    value = ((self.step.numpy() - self.last_iteration) /
             (time.time() - self.last_time))
    if restart:
      self.start()
    return value