standard_runner.py 14.1 KB
Newer Older
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
# Copyright 2020 The Orbit Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     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.
# ==============================================================================
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
"""AbstractTrainer/Evaluator subclasses with added functionality.

The classes in this module provide some additional structure to the bare
`AbstractTrainer`/`AbstractEvaluator` APIs.

Both `StandardTrainer` and `StandardEvaluator` split the train/eval loops into
"begin", "step", and "end" methods, and provide an implementation of the loop
itself that makes calls to the relevant step method.

`StandardTrainer` supports running the loop using the TF while loop construct
for added performance (particularly on TPUs). It additionally provides some
functionality to make writing summaries from inside a model more performant when
running on TPUs.

These classes are intended to work well in common settings, however there may
be use cases these classes don't support (for instance, `StandardEvaluator` in
particular doesn't support running full evaluations over multiple different eval
datasets). Users are encouraged to simply fall back to custom `AbstractTrainer`
and `AbstractEvaluator` subclasses in these cases.
"""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
35
36

import abc
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
37

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
38
from typing import Any, Optional
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
39

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
40
import dataclasses
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
41

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
42
from orbit import runner
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
43
44
from orbit.utils import loop_fns

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
45
46
47
import tensorflow as tf


A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
48
@dataclasses.dataclass(frozen=True)
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
49
50
class StandardTrainerOptions:
  """Advanced options for `orbit.StandardTrainer`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
51
52

  Attributes:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
53
54
55
56
57
58
59
60
61
62
63
64
    use_tf_while_loop: A boolean indicating whether to run the training loop
      using a `tf.while_loop`. If `True`, `use_tf_function` must also be `True`.
    use_tf_function: A boolean indicating whether to apply `tf.function` to the
      training loop. This will only affect the body of the loop (involving
      `train_step`); `train_loop_begin` and `train_loop_end` will always be run
      in eager mode.
    use_tpu_summary_optimization: A boolean indicating whether to enable a
      performance optimization for summaries in TPUs. Writing summaries
      conditionally with outside compilation on TPUs can be extremely slow. If
      `True`, this optimization creates two `tf.function`s with two XLA programs
      (one with summary calls, and one without). The program with summaries runs
      only for one step when summaries should be recorded.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
65
66
67
68
69
70
  """
  use_tf_while_loop: bool = True
  use_tf_function: bool = True
  use_tpu_summary_optimization: bool = False


Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def _create_train_loop_fn(train_step_fn, options: StandardTrainerOptions):
  """Creates a training loop from the given step function and options."""
  if options.use_tf_while_loop:
    loop_fn = loop_fns.create_tf_while_loop_fn(train_step_fn)
    if options.use_tpu_summary_optimization:
      loop_fn = loop_fns.LoopFnWithSummaries(loop_fn)
    else:
      loop_fn = tf.function(loop_fn)
  else:
    if options.use_tf_function:
      train_step_fn = tf.function(train_step_fn)
    loop_fn = loop_fns.create_loop_fn(train_step_fn)
  return loop_fn


Hongkun Yu's avatar
Hongkun Yu committed
86
class StandardTrainer(runner.AbstractTrainer, metaclass=abc.ABCMeta):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
87
88
89
90
91
92
93
94
95
96
97
98
99
  """Implements standard functionality on top of the AbstractTrainer API.

  This class structures the training "inner loop" roughly as follows:

      train_loop_begin()
      for _ in range(num_steps):
        train_step(train_iterator)
      return train_loop_end()

  Calls to `train_loop_begin` and `train_loop_end` are always done in eager
  mode, while the loop/`train_step` may be implemented using `tf.while` and/or
  `tf.function`, as determined by the `options` passed to `__init__`.
  """
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
100

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
101
  def __init__(self, train_dataset, options: StandardTrainerOptions = None):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
102
    """Initializes the `StandardTrainer` instance.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
103
104

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
105
106
      train_dataset: A `tf.nest`-compatible structure of `tf.data.Dataset` or
        `DistributedDataset`.
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
107
      options: An `orbit.StandardTrainerOptions` instance.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
108
    """
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
109
110
    options = options or StandardTrainerOptions()
    if options.use_tf_while_loop and not options.use_tf_function:
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
111
112
      raise ValueError("`use_tf_while_loop=True` and `use_tf_function=False` "
                       "is not supported")
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
113
    if options.use_tpu_summary_optimization and not options.use_tf_while_loop:
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
114
115
      raise ValueError("`use_tpu_summary_optimization=True` and "
                       "`use_tf_while_loop=False` is not supported")
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
116

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
117
    self._train_options = options
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
118
119
120
121
    self._train_dataset = train_dataset
    self._train_iter = None
    self._train_loop_fn = None

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
122
123
124
125
126
127
128
129
130
131
  def train(self, num_steps: tf.Tensor) -> Optional[runner.Output]:
    """Implements `num_steps` steps of training.

    Args:
      num_steps: The number of training steps to run. This corresponds directly
        to the number of calls made to `train_step`.

    Returns:
      The output of `train_loop_end`.
    """
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
132
133
    self.train_loop_begin()

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
134
135
136
137
    if self._train_loop_fn is None:
      self._train_loop_fn = _create_train_loop_fn(
          self.train_step, options=self._train_options)

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
138
139
140
141
142
143
144
145
146
    if self._train_iter is None:
      self._train_iter = tf.nest.map_structure(iter, self.train_dataset)

    self._train_loop_fn(self._train_iter, num_steps)
    return self.train_loop_end()

  def train_loop_begin(self):
    """Called once at the beginning of the training loop.

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
147
148
149
150
    This method is always called in eager mode, and is a good place to reset
    metrics that accumulate values over multiple steps of training.

    Note that this method is called before dataset iterator creation.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
151
152
153
154
155
156
157
    """
    pass

  @abc.abstractmethod
  def train_step(self, iterator):
    """Implements one step of training.

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
158
159
    What a "step" consists of is up to the implementer. When using distribution
    strategies, the call to this method takes place in the "cross-replica
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
160
161
162
    context" for generality, to allow e.g. multiple iterator dequeues and calls
    to `strategy.run`.

Ruoxin Sang's avatar
Ruoxin Sang committed
163
    Note that if `use_tf_function=True`, all the code inside `train_step` should
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
164
165
166
167
    be compatible with `tf.function` tracing (and in particular, any state
    modifications involving `self` should be avoided). In some cases, non-
    `tf.function` compatible code can be moved to `train_loop_begin` or
    `train_loop_end`, which always execute eagerly.
Ruoxin Sang's avatar
Ruoxin Sang committed
168

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
169
    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
170
171
172
      iterator: A `tf.nest`-compatible structure of `tf.data.Iterator` or
        `DistributedIterator`. The structure of this input matches the structure
        of `train_dataset` as passed to `__init__`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
173
174
175
    """
    pass

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
176
177
  def train_loop_end(self) -> Optional[runner.Output]:
    """Called once at the end of the training loop.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
178

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
179
180
181
    This method is always called in eager mode, and is a good place to get
    metric results. The value returned from this function will be returned as-is
    from the `train` method implementation provided by `StandardTrainer`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
182
183
184

    Returns:
      The function may return a dictionary of `Tensors`, which will be
Ruoxin Sang's avatar
Ruoxin Sang committed
185
186
      written to logs and as TensorBoard summaries. It can also be a
      nested dictionary, yielding a hierarchy of summary directories.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
187
188
189
190
191
    """
    pass

  @property
  def train_dataset(self):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
192
    """The current training dataset."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
193
194
195
196
    return self._train_dataset

  @train_dataset.setter
  def train_dataset(self, train_dataset):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
197
    """Sets a new training dataset, replacing the current one.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
198

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
199
    Any unprocessed examples in the current dataset are discarded.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
200
201

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
202
203
      train_dataset: A `tf.nest`-compatible structure of `tf.data.Dataset` or
        `DistributedDataset`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
204
205
206
207
208
    """
    self._train_dataset = train_dataset
    self._train_iter = None


A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
209
@dataclasses.dataclass(frozen=True)
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
210
211
class StandardEvaluatorOptions:
  """Advanced options for the `orbit.StandardEvaluator`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
212
213

  Attributes:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
214
215
216
217
    use_tf_function: A boolean indicating whether to apply `tf.function` to the
      training loop. This will only affect the body of the loop (involving
      `train_step`); `train_loop_begin` and `train_loop_end` will always be run
      in eager mode.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
218
  """
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
219
  use_tf_function: bool = True
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
220
221


Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
222
223
224
225
226
227
def _create_eval_loop_fn(eval_step_fn, options: StandardEvaluatorOptions):
  if options.use_tf_function:
    eval_step_fn = tf.function(eval_step_fn)
  return loop_fns.create_loop_fn(eval_step_fn)


Hongkun Yu's avatar
Hongkun Yu committed
228
class StandardEvaluator(runner.AbstractEvaluator, metaclass=abc.ABCMeta):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
  """Implements the standard functionality of AbstractEvaluator APIs.

  This class structures evaluation roughly as follows:

      state = eval_begin()
      for _ in range(num_steps):
        step_outputs = eval_step(eval_iterator)
        state = eval_reduce(state, step_outputs)
      return eval_end(state)

  Calls to `eval_begin`, `eval_reduce`, and `eval_end` are always done in eager
  mode, while `eval_step` may be compiled with `tf.function` as determined by
  the `options` passed to `__init__`.

  This class does not support completely evaluating multiple different datasets
  (i.e., where every example of each dataset should be processed, as opposed to
  running for a fixed number of evaluation steps). A custom `AbstractEvaluator`
  is recommended in this case.
  """
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
248

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
249
  def __init__(self, eval_dataset, options: StandardEvaluatorOptions = None):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
250
    """Initializes the `StandardEvaluator` instance.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
251
252

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
253
254
      eval_dataset: A `tf.nest`-compatible structure of `tf.data.Dataset` or
        `DistributedDataset`.
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
255
      options: An `orbit.StandardEvaluatorOptions` instance.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
256
    """
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
257
    self._eval_options = options or StandardEvaluatorOptions()
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
258
259
260
    self._eval_dataset = eval_dataset
    self._eval_loop_fn = None

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
261
262
263
264
265
266
267
268
269
270
271
  def evaluate(self, num_steps: tf.Tensor) -> Optional[runner.Output]:
    """Implements `num_steps` steps of evaluation.

    Args:
      num_steps: The number of evaluation steps to run. When this is -1,
        evaluation proceeds until a call to `eval_step` raises a `StopIteration`
        or `tf.errors.OutOfRangeError`.

    Returns:
      The output of `self.eval_end()`.
    """
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
272
273
274
    outputs = self.eval_begin()  # pylint: disable=assignment-from-no-return

    if self._eval_loop_fn is None:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
275
276
      self._eval_loop_fn = _create_eval_loop_fn(
          self.eval_step, options=self._eval_options)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
277

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
278
    eval_iter = tf.nest.map_structure(iter, self.eval_dataset)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
279
280
    outputs = self._eval_loop_fn(
        eval_iter, num_steps, state=outputs, reduce_fn=self.eval_reduce)
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
281

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
282
283
284
285
286
287
288
289
    if outputs is None:
      return self.eval_end()
    else:
      return self.eval_end(outputs)

  def eval_begin(self) -> Any:
    """Called once at the beginning of the evaluation.

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
290
291
292
293
    This method is always called in eager mode, and is a good place to reset
    metrics that accumulate values over the course of evaluation.

    Note that this method is called before dataset iterator creation.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
294
295

    Returns:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
296
      An value to pass as the `state` argument to `eval_reduce`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
297
298
299
300
301
302
303
    """
    pass

  @abc.abstractmethod
  def eval_step(self, iterator) -> Any:
    """Implements one step of evaluation.

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
304
305
    What a "step" consists of is up to the implementer. When using distribution
    strategies, the call to this method takes place in the "cross-replica
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
306
307
308
    context" for generality, to allow e.g. multiple iterator dequeues and calls
    to `strategy.run`.

Ruoxin Sang's avatar
Ruoxin Sang committed
309
    Note that if `use_tf_function=True`, all the code inside `eval_step` should
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
310
311
312
313
    be compatible with `tf.function` tracing (and in particular, any state
    modifications involving `self` should be avoided). In some cases, non-
    `tf.function` compatible code can be moved to `eval_loop_begin`,
    `eval_reduce`, or `eval_loop_end`, which always execute eagerly.
Ruoxin Sang's avatar
Ruoxin Sang committed
314

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
315
    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
316
317
      iterator: A `tf.nest`-compatible structure of `tf.data.Iterator` or
        `DistributedIterator`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
318
319
320
321
322
323
324

    Returns:
      An output which is passed as `step_outputs` argument into `eval_reduce`
      function.
    """
    pass

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
325
  def eval_end(self, *args) -> Optional[runner.Output]:
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
326
327
    """Called at the end of the evaluation.

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
328
329
330
331
332
    Called once at the end of evaluation.

    This method is always called in eager mode, and is a good place to get
    metric results. The value returned from this function will be returned as-is
    from the `evaluate` method implementation provided by `StandardEvaluator`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
333
334

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
335
336
      *args: The outputs from `eval_reduce` for the last eval step, if they are
        non-`None` (if they are `None`, nothing is passed).
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
337
338
339

    Returns:
      The function may return a dictionary of `Tensors`, which will be
Ruoxin Sang's avatar
Ruoxin Sang committed
340
341
      written to logs and as TensorBoard summaries. It can also be a
      nested dictionary, yielding a hierarchy of summary directories.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
342
343
344
    """
    pass

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
345
346
347
348
  def eval_reduce(self,
                  state: Any = None,
                  step_outputs: Optional[runner.Output] = None) -> Any:
    """A function to perform per-step reduction on the evaluation outputs.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
349

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
350
351
352
353
354
    This is useful for passing state throughout evaluation, especially in cases
    where maintaining or accumulating state is hard to accomplish using
    `tf.metrics.Metric` or other `tf.Variable`-based approaches. For instance,
    it can be used to easily accumulate all per-example losses from the full
    evaluation for subsequent processing in `eval_end()`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
355
356

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
357
      state: A state being mainted throughout the evaluation.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
358
359
360
      step_outputs: Outputs from the current evaluation step.

    Returns:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
361
362
363
      An output which is passed as the `state` argument to this function for the
      next step. After evaluation is finished, the output from last step will be
      passed to `eval_end`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
364
365
366
367
368
    """
    pass

  @property
  def eval_dataset(self):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
369
    """The current evaluation dataset."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
370
371
372
373
    return self._eval_dataset

  @eval_dataset.setter
  def eval_dataset(self, eval_dataset):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
374
375
376
    """Sets a new eval dataset, replacing the current one.

    Any unprocessed examples in the current dataset are discarded.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
377
378

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
379
380
      eval_dataset: A `tf.nest`-compatible structure of `tf.data.Dataset` or
        `DistributedDataset`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
381
382
    """
    self._eval_dataset = eval_dataset