controller.py 21.8 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

Ron Shapiro's avatar
Ron Shapiro committed
20
from typing import 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
104
      eval_summary_dir: Optional[str] = None,
  ):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
105
106
107
108
109
    """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
110
111

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
112
113
114
115
116
117
118
119
120
      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()`).
121
122
123
124
125
126
127
      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
128
129
130
131
132
      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`.
133
134
135
136
137
      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
138
139
140
141
142
      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
143
      summary_interval: Step interval for training summaries. Note that this
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
144
145
146
147
148
149
150
151
152
153
154
        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.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
155
156

    Raises:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
157
      ValueError: If both `trainer` and `evaluator` are `None`.
158
      ValueError: If `steps_per_loop` is not a positive integer or a callable.
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
159
160
      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
161
162
    """
    if trainer is None and evaluator is None:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
163
      raise ValueError("`trainer` and `evaluator` should not both be `None`.")
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
164

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
165
166
    if trainer is not None:
      if steps_per_loop is None:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
167
168
        raise ValueError(
            "`steps_per_loop` is required when `trainer` is provided.")
169
170
      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
171
        raise ValueError(
172
173
            f"`steps_per_loop` ({steps_per_loop}) must be a positive integer "
            "or a callable.")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
174
175
176

      if summary_interval is not None:
        if summary_interval <= 0:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
177
178
          raise ValueError(
              f"`summary_interval` ({summary_interval}) must be larger than 0.")
179
180
        elif not callable(steps_per_loop) and (summary_interval % steps_per_loop
                                               != 0):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
181
182
183
184
          raise ValueError(
              f"`summary interval` ({summary_interval}) must be a multiple "
              f"of `steps_per_loop` ({steps_per_loop}).")

185
    if not isinstance(global_step, tf.Variable):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
186
      raise ValueError("`global_step` must be a `tf.Variable`.")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
187
188
189
190
191
192

    self.trainer = trainer
    self.evaluator = evaluator

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

Ron Shapiro's avatar
Ron Shapiro committed
193
194
    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
195

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
196
197
198
199
200
201
202
203
    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
      self.summary_manager = utils.SummaryManager(
          summary_dir, tf.summary.scalar, global_step=self.global_step)
204
      self._steps_per_loop = steps_per_loop
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
205
206
207
208
209
210
211
212
213
214
215

    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:
        self.eval_summary_manager = utils.SummaryManager(
            eval_summary_dir, tf.summary.scalar, global_step=self.global_step)

216
    tf.summary.experimental.set_step(self.global_step)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
217
218
219

    # Restores the model if needed.
    if self.checkpoint_manager is not None:
220
221
      restored_path = self.restore_checkpoint()
      if restored_path:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
222
        _log(f"restored from checkpoint: {restored_path}")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
223
224

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

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
227
228
229
230
    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
231
232
233
234

    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
235
        returns (regardless of the checkpointing interval). Defaults to `True`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
236
    """
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
237
    self._require("trainer", for_method="train")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
238
239

    # TODO(momernick): Support steps=None or -1 (training to exhaustion).
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
240
241
    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
242
243
244
245
246
    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
247
      current_step = self.global_step.numpy()
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
248
249

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

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
252
253
  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
254

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
255
256
    This method calls `self.evaluator.evaluate(steps)`, then writes the returned
    summaries (if any).
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
257
258

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
259
260
261
262
      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
263

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

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
267
    Raises:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
268
269
270
      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
271
    """
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
272
    self._require("evaluator", for_method="evaluate")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
273
274

    if steps > 0:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
275
276
277
      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
278
    else:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
279
      raise ValueError(f"`steps` ({steps}) should be > 0, or == -1.")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
280

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

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
284
285
286
287
288
    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
289

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
290
291
292
293
294
    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
295
    _log(f" eval | step: {current_step: 6d} | "
296
         f"eval time: {elapsed: 6.1f} sec | "
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
297
         f"output: {_format_output(eval_output)}")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
298

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
299
    self.eval_summary_manager.write_summaries(eval_output)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
300
301
    self.eval_summary_manager.flush()

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
302
    return eval_output
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
303
304

  def train_and_evaluate(self,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
305
                         train_steps: int,
306
                         eval_steps: int = -1,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
307
                         eval_interval: Optional[int] = None) -> None:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
308
    """Runs interleaved training and evaluation.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
309

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
310
311
312
313
314
    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
315
316
317

    Args:
      train_steps: The global step count to train up to.
318
      eval_steps: The number of steps to run during an evaluation. If -1, this
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
319
320
321
322
        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
323
324
        setting. If None, evaluation will only be performed after training is
        complete.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
325
    """
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
326
327
328
329
    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
330
331
332
333
334
335
    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
336
337
      current_step = self.global_step.numpy()
    self._maybe_save_checkpoint(check_interval=False)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
338
339

  def evaluate_continuously(self,
340
                            steps: int = -1,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
341
342
                            timeout: Optional[Union[int, float]] = None,
                            timeout_fn: Optional[Callable[[], bool]] = None):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
343
    """Continuously monitors a directory and evaluates new checkpoints in it.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
344
345
346
347
348
349

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

    Args:
350
351
      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
352
353
354
355
356
357
358
359
360
361
      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
362
363
364
    self._require("evaluator", for_method="evaluate_continuously")
    self._require("checkpoint_manager", for_method="evaluate_continuously")

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
365
366
367
368
369
370
371
    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
372
  def restore_checkpoint(self, checkpoint_path: Optional[str] = None):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
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
    """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)

416
417
418
419
420
421
422
  @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
423
  def _train_n_steps(self, num_steps: int):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
424
    """Runs training for `num_steps` steps.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
425

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
426
427
428
    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
429
430

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

    Raises:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
434
435
      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
436
437
438
439
440
    """
    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
441
    with self.summary_manager.summary_writer().as_default():
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
442
443
      should_record = False  # Allows static optimization in no-summary cases.
      if self.summary_interval:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
444
        # Create a predicate to determine when summaries should be written.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
445
446
        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
447
448
449
450
451
452
        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:
453
      message = (
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
454
455
456
          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()}.")
457
      logging.warning(message)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
458

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
459
460
461
462
463
    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)

464
    current_step = self.global_step.numpy()
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
465
    steps_per_second = self.step_timer.steps_per_second()
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
466
467
468
469
470
471
472
    _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
473

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

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
477
478
479
    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
480
481

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
482
483
484
485
      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
486
487
488
489
490
491
492

    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
493
          check_interval=check_interval)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
494
      if ckpt_path is not None:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
495
        _log(f"saved checkpoint to {ckpt_path}.")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
496
497
498
        return True
    return False

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
499
500
501
502
503
504
505
  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
506

Hongkun Yu's avatar
Hongkun Yu committed
507
class StepTimer:
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
  """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