standard_runner.py 16.5 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
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
    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.
57
58
    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`.
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
59
60
61
62
63
64
    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
  """
  use_tf_function: bool = True
67
  use_tf_while_loop: bool = True
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
68
69
70
  use_tpu_summary_optimization: bool = False


Hongkun Yu's avatar
Hongkun Yu committed
71
class StandardTrainer(runner.AbstractTrainer, metaclass=abc.ABCMeta):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
72
73
74
75
76
77
78
79
80
81
82
83
84
  """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
85

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
86
  def __init__(self, train_dataset, options: StandardTrainerOptions = None):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
87
    """Initializes the `StandardTrainer` instance.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
88
89

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
90
91
      train_dataset: A `tf.nest`-compatible structure of `tf.data.Dataset` or
        `DistributedDataset`.
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
92
      options: An `orbit.StandardTrainerOptions` instance.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
93
    """
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
94
95
    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
96
97
      raise ValueError("`use_tf_while_loop=True` and `use_tf_function=False` "
                       "is not supported")
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
98
    if options.use_tpu_summary_optimization and not options.use_tf_while_loop:
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
99
100
      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
101

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
102
    self._train_options = options
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
103
104
105
106
    self._train_dataset = train_dataset
    self._train_iter = None
    self._train_loop_fn = None

107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
  def create_train_loop_fn(self):
    """Creates a training loop from the current step function and options.

    Returns:
      The train loop function, i.e. wrapper of multiple train steps.
    """
    train_step_fn = self.train_step
    if self._train_options.use_tf_while_loop:
      loop_fn = loop_fns.create_tf_while_loop_fn(train_step_fn)
      if self._train_options.use_tpu_summary_optimization:
        loop_fn = loop_fns.LoopFnWithSummaries(loop_fn)
      else:
        loop_fn = tf.function(loop_fn)
    else:
      if self._train_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

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
126
127
128
129
130
131
132
133
134
135
  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
136
137
    self.train_loop_begin()

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
138
    if self._train_loop_fn is None:
139
      self._train_loop_fn = self.create_train_loop_fn()
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
140

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
141
142
143
144
145
146
147
148
149
    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
150
151
152
153
    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
154
155
156
157
158
159
160
    """
    pass

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

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
161
162
    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
163
164
165
    context" for generality, to allow e.g. multiple iterator dequeues and calls
    to `strategy.run`.

Ruoxin Sang's avatar
Ruoxin Sang committed
166
    Note that if `use_tf_function=True`, all the code inside `train_step` should
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
167
168
169
170
    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
171

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
172
    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
173
174
175
      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
176
177
178
    """
    pass

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
179
180
  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
181

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
182
183
184
    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
185
186
187

    Returns:
      The function may return a dictionary of `Tensors`, which will be
Ruoxin Sang's avatar
Ruoxin Sang committed
188
189
      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
190
191
192
193
194
    """
    pass

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

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

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
202
    Any unprocessed examples in the current dataset are discarded.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
203
204

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


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

  Attributes:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
217
218
219
220
    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.
221
222
    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`.
223
224
225
226
227
228
229
    recreate_iterator_for_each_eval: A boolean indicating whether to recreate a
      new iterator for the evaluation dataset before each round of evaluation,
      which implies each round of evaluation starts from the beginning of
      the evaluation dataset. For example, the evaluation dataset is
      `[1, 2, 3, 4]`, batch size is 1 and evaluation steps is 2. If `True`, the
      data to be evaluated is [1, 2] every time. If `False`, the iterator
      state is maintained between calls to `StandardEvaluator.evaluate()`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
230
  """
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
231
  use_tf_function: bool = True
232
  use_tf_while_loop: bool = False
233
  recreate_iterator_for_each_eval: bool = True
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
234
235


Hongkun Yu's avatar
Hongkun Yu committed
236
class StandardEvaluator(runner.AbstractEvaluator, metaclass=abc.ABCMeta):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
  """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
256

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
257
  def __init__(self, eval_dataset, options: StandardEvaluatorOptions = None):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
258
    """Initializes the `StandardEvaluator` instance.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
259
260

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
261
262
      eval_dataset: A `tf.nest`-compatible structure of `tf.data.Dataset` or
        `DistributedDataset`.
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
263
      options: An `orbit.StandardEvaluatorOptions` instance.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
264
    """
265
266
267
268
269
270
    options = options or StandardEvaluatorOptions()
    if options.use_tf_while_loop and not options.use_tf_function:
      raise ValueError("`use_tf_while_loop=True` and `use_tf_function=False` "
                       "is not supported")

    self._eval_options = options
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
271
    self._eval_dataset = eval_dataset
272
    self._eval_iter = None
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
273
274
    self._eval_loop_fn = None

275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
  def create_eval_loop_fn(self, has_state: bool):
    """Creates an eval loop from the current step function and options.

    Args:
      has_state: If the step function has state, state will be kept in the loop.

    Returns:
      The eval loop function, i.e. wrapper of multiple eval steps.
    """
    eval_step_fn = self.eval_step
    if self._eval_options.use_tf_while_loop:
      # TODO(b/176126742): tf.while_loop doesn't support `None` as a loop input
      # even when it is not used inside the loop. To workaround this limitation,
      # we have to build two tf.functions for it.
      if has_state:
        loop_fn = loop_fns.create_tf_while_loop_fn_with_state(eval_step_fn)
      else:
        loop_fn = loop_fns.create_tf_while_loop_fn(eval_step_fn)
      loop_fn = tf.function(loop_fn)
    else:
      if self._eval_options.use_tf_function:
        eval_step_fn = tf.function(eval_step_fn)
      loop_fn = loop_fns.create_loop_fn(eval_step_fn)
    return loop_fn

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
300
301
302
303
304
305
306
307
308
309
  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()`.
310
311
312
313

    Raises:
      ValueError: If `options.use_tf_while_loop` is `True` and `num_steps` is
        unspecified.
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
314
    """
315
316
317
318
    if self._eval_options.use_tf_while_loop and num_steps == -1:
      raise ValueError("Looping until exhausted is not supported if "
                       "`options.use_tf_while_loop` is `True`")

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
319
320
    outputs = self.eval_begin()  # pylint: disable=assignment-from-no-return

321
    has_state = outputs is not None
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
322
    if self._eval_loop_fn is None:
323
      self._eval_loop_fn = self.create_eval_loop_fn(has_state)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
324

325
326
327
328
329
330
331
332
333
    # If `recreate_iterator_for_each_eval` is `True`, `self._eval_iter` is
    # always None.
    if self._eval_iter is None:
      eval_iter = tf.nest.map_structure(iter, self.eval_dataset)
      if not self._eval_options.recreate_iterator_for_each_eval:
        self._eval_iter = eval_iter
    else:
      eval_iter = self._eval_iter

334
335
336
337
338
    if self._eval_options.use_tf_while_loop and not has_state:
      self._eval_loop_fn(eval_iter, num_steps)
    else:
      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
339

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
340
341
342
343
344
345
346
347
    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
348
349
350
351
    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
352
353

    Returns:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
354
      An value to pass as the `state` argument to `eval_reduce`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
355
356
357
358
359
360
361
    """
    pass

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

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
362
363
    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
364
365
366
    context" for generality, to allow e.g. multiple iterator dequeues and calls
    to `strategy.run`.

Ruoxin Sang's avatar
Ruoxin Sang committed
367
    Note that if `use_tf_function=True`, all the code inside `eval_step` should
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
368
369
370
371
    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
372

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
373
    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
374
375
      iterator: A `tf.nest`-compatible structure of `tf.data.Iterator` or
        `DistributedIterator`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
376
377
378
379
380
381
382

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

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

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

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
393
394
      *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
395
396
397

    Returns:
      The function may return a dictionary of `Tensors`, which will be
Ruoxin Sang's avatar
Ruoxin Sang committed
398
399
      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
400
401
402
    """
    pass

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
403
404
405
406
  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
407

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
408
409
410
411
412
    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
413
414

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
415
      state: A state being mainted throughout the evaluation.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
416
417
418
      step_outputs: Outputs from the current evaluation step.

    Returns:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
419
420
421
      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
422
423
424
425
426
    """
    pass

  @property
  def eval_dataset(self):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
427
    """The current evaluation dataset."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
428
429
430
431
    return self._eval_dataset

  @eval_dataset.setter
  def eval_dataset(self, eval_dataset):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
432
433
434
    """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
435
436

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
437
438
      eval_dataset: A `tf.nest`-compatible structure of `tf.data.Dataset` or
        `DistributedDataset`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
439
440
    """
    self._eval_dataset = eval_dataset
441
    self._eval_iter = None