controller.py 19.9 KB
Newer Older
1
# Copyright 2021 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

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
20
from typing import Callable, 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)


Hongkun Yu's avatar
Hongkun Yu committed
49
class Controller:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
  """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`,
  and interleaves calls to the user provided inner loops with additional actions
  such as saving checkpoints, running evaluations, and writing summaries
  (depending on the arguments passed to `Controller.__init__` and the method
  being called).

  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.
  """
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
74
75
76

  def __init__(
      self,
77
78
      *,  # Makes all args keyword only.
      global_step: tf.Variable,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
79
80
      trainer: Optional[runner.AbstractTrainer] = None,
      evaluator: Optional[runner.AbstractEvaluator] = None,
81
      strategy: Optional[tf.distribute.Strategy] = None,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
82
83
84
85
86
      # Train related
      steps_per_loop: Optional[int] = None,
      checkpoint_manager: Optional[tf.train.CheckpointManager] = None,
      # Summary related
      summary_interval: Optional[int] = None,
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
87
      summary_dir: Optional[str] = None,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
88
      # Evaluation related
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
89
90
91
92
93
94
      eval_summary_dir: Optional[str] = None):
    """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
95
96

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
97
98
99
100
101
102
103
104
105
      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()`).
106
107
108
109
110
111
112
      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()`.
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
113
114
115
116
117
118
119
      steps_per_loop: The number of steps to run in each inner loop of training
        (passed as the `num_steps` parameter of `trainer.train`).
      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
120
      summary_interval: Step interval for training summaries. Note that this
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
121
122
123
124
125
126
127
128
129
130
131
        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
132
133

    Raises:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
134
      ValueError: If both `trainer` and `evaluator` are `None`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
135
      ValueError: If `steps_per_loop` is not a positive integer.
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
136
137
      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
138
139
    """
    if trainer is None and evaluator is None:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
140
      raise ValueError("`trainer` and `evaluator` should not both be `None`.")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
141
142
    if trainer is not None:
      if steps_per_loop is None:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
143
144
145
146
147
        raise ValueError(
            "`steps_per_loop` is required when `trainer` is provided.")
      elif not isinstance(steps_per_loop, int) or steps_per_loop < 1:
        raise ValueError(
            f"`steps_per_loop` ({steps_per_loop}) must be a positive integer.")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
148
149
150

      if summary_interval is not None:
        if summary_interval <= 0:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
151
152
153
154
155
156
157
          raise ValueError(
              f"`summary_interval` ({summary_interval}) must be larger than 0.")
        elif summary_interval % steps_per_loop != 0:
          raise ValueError(
              f"`summary interval` ({summary_interval}) must be a multiple "
              f"of `steps_per_loop` ({steps_per_loop}).")

158
    if not isinstance(global_step, tf.Variable):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
159
      raise ValueError("`global_step` must be a `tf.Variable`.")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185

    self.trainer = trainer
    self.evaluator = evaluator

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

    self.global_step = global_step
    self.checkpoint_manager = checkpoint_manager

    if self.trainer is not None:
      self.step_timer = None
      self.steps_per_loop = steps_per_loop
      self.summary_interval = summary_interval
      self.summary_manager = utils.SummaryManager(
          summary_dir, tf.summary.scalar, global_step=self.global_step)

    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)

186
    tf.summary.experimental.set_step(self.global_step)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
187
188
189

    # Restores the model if needed.
    if self.checkpoint_manager is not None:
190
191
      restored_path = self.restore_checkpoint()
      if restored_path:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
192
        _log(f"restored from checkpoint: {restored_path}")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
193
194

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

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
197
198
199
200
    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
201
202
203
204

    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
205
        returns (regardless of the checkpointing interval). Defaults to `True`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
206
    """
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
207
    self._require("trainer", for_method="train")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
208
209

    # TODO(momernick): Support steps=None or -1 (training to exhaustion).
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
210
211
    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
212
213
214
215
216
    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
217
      current_step = self.global_step.numpy()
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
218
219

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

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
222
223
  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
224

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
225
226
    This method calls `self.evaluator.evaluate(steps)`, then writes the returned
    summaries (if any).
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
227
228

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

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

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
237
    Raises:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
238
239
240
      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
241
    """
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
242
    self._require("evaluator", for_method="evaluate")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
243
244

    if steps > 0:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
245
246
247
      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
248
    else:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
249
      raise ValueError(f"`steps` ({steps}) should be > 0, or == -1.")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
250

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

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
254
255
256
257
258
259
    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)
    eval_output = tf.nest.map_structure(utils.get_value, eval_output or {})
    elapsed = time.time() - start
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
260

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

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
265
    self.eval_summary_manager.write_summaries(eval_output)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
266
267
    self.eval_summary_manager.flush()

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
268
    return eval_output
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
269
270

  def train_and_evaluate(self,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
271
                         train_steps: int,
272
                         eval_steps: int = -1,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
273
                         eval_interval: Optional[int] = None) -> None:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
274
    """Runs interleaved training and evaluation.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
275

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
276
277
278
279
280
    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
281
282
283

    Args:
      train_steps: The global step count to train up to.
284
      eval_steps: The number of steps to run during an evaluation. If -1, this
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
285
286
287
288
        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
289
290
        setting. If None, evaluation will only be performed after training is
        complete.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
291
292
293
294

    Raises:
      ValueError: If eval_interval is not a multiple of self.steps_per_loop.
    """
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
295
296
297
298
    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
299
300
301
302
303
304
    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
305
306
      current_step = self.global_step.numpy()
    self._maybe_save_checkpoint(check_interval=False)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
307
308

  def evaluate_continuously(self,
309
                            steps: int = -1,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
310
311
                            timeout: Optional[Union[int, float]] = None,
                            timeout_fn: Optional[Callable[[], bool]] = None):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
312
    """Continuously monitors a directory and evaluates new checkpoints in it.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
313
314
315
316
317
318

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

    Args:
319
320
      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
321
322
323
324
325
326
327
328
329
330
      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
331
332
333
    self._require("evaluator", for_method="evaluate_continuously")
    self._require("checkpoint_manager", for_method="evaluate_continuously")

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
334
335
336
337
338
339
340
    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
341
  def restore_checkpoint(self, checkpoint_path: Optional[str] = None):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
    """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)

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

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
388
389
390
    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
391
392

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

    Raises:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
396
397
      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
398
399
400
401
402
    """
    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
403
    with self.summary_manager.summary_writer().as_default():
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
404
405
      should_record = False  # Allows static optimization in no-summary cases.
      if self.summary_interval:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
406
        # Create a predicate to determine when summaries should be written.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
407
408
        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
409
410
411
412
413
414
415
        num_steps_tensor = tf.convert_to_tensor(num_steps, dtype=tf.int32)
        train_output = self.trainer.train(num_steps_tensor)
    train_output = tf.nest.map_structure(utils.get_value, train_output or {})

    # Verify that global_step was updated properly, then update current_step.
    expected_step = current_step + num_steps
    if self.global_step.numpy() != expected_step:
416
      message = (
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
417
418
419
          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()}.")
420
421
      logging.warning(message)
      return
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
422

423
    current_step = expected_step
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
424
    steps_per_second = self.step_timer.steps_per_second()
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
425
426
427
428
429
430
431
    _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
432

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

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
436
437
438
    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
439
440

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
441
442
443
444
      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
445
446
447
448
449
450
451

    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
452
          check_interval=check_interval)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
453
      if ckpt_path is not None:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
454
        _log(f"saved checkpoint to {ckpt_path}.")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
455
456
457
        return True
    return False

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

Hongkun Yu's avatar
Hongkun Yu committed
466
class StepTimer:
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
  """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