standard_runner.py 15.4 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
# 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.
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


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.
218
219
    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`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
220
  """
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
221
  use_tf_function: bool = True
222
  use_tf_while_loop: bool = False
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
223
224


225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
def _create_eval_loop_fn(eval_step_fn, has_state: bool,
                         options: StandardEvaluatorOptions):
  """Create evaluation loop function."""
  if 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 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
242
243


Hongkun Yu's avatar
Hongkun Yu committed
244
class StandardEvaluator(runner.AbstractEvaluator, metaclass=abc.ABCMeta):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
  """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
264

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
265
  def __init__(self, eval_dataset, options: StandardEvaluatorOptions = None):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
266
    """Initializes the `StandardEvaluator` instance.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
267
268

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
269
270
      eval_dataset: A `tf.nest`-compatible structure of `tf.data.Dataset` or
        `DistributedDataset`.
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
271
      options: An `orbit.StandardEvaluatorOptions` instance.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
272
    """
273
274
275
276
277
278
    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
279
280
281
    self._eval_dataset = eval_dataset
    self._eval_loop_fn = None

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
282
283
284
285
286
287
288
289
290
291
  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()`.
292
293
294
295

    Raises:
      ValueError: If `options.use_tf_while_loop` is `True` and `num_steps` is
        unspecified.
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
296
    """
297
298
299
300
    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
301
302
    outputs = self.eval_begin()  # pylint: disable=assignment-from-no-return

303
    has_state = outputs is not None
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
304
    if self._eval_loop_fn is None:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
305
      self._eval_loop_fn = _create_eval_loop_fn(
306
          self.eval_step, has_state=has_state, options=self._eval_options)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
307

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
308
    eval_iter = tf.nest.map_structure(iter, self.eval_dataset)
309
310
311
312
313
    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
314

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
315
316
317
318
319
320
321
322
    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
323
324
325
326
    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
327
328

    Returns:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
329
      An value to pass as the `state` argument to `eval_reduce`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
330
331
332
333
334
335
336
    """
    pass

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

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
337
338
    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
339
340
341
    context" for generality, to allow e.g. multiple iterator dequeues and calls
    to `strategy.run`.

Ruoxin Sang's avatar
Ruoxin Sang committed
342
    Note that if `use_tf_function=True`, all the code inside `eval_step` should
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
343
344
345
346
    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
347

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
348
    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
349
350
      iterator: A `tf.nest`-compatible structure of `tf.data.Iterator` or
        `DistributedIterator`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
351
352
353
354
355
356
357

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

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

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
361
362
363
364
365
    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
366
367

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
368
369
      *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
370
371
372

    Returns:
      The function may return a dictionary of `Tensors`, which will be
Ruoxin Sang's avatar
Ruoxin Sang committed
373
374
      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
375
376
377
    """
    pass

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
378
379
380
381
  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
382

Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
383
384
385
386
387
    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
388
389

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
390
      state: A state being mainted throughout the evaluation.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
391
392
393
      step_outputs: Outputs from the current evaluation step.

    Returns:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
394
395
396
      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
397
398
399
400
401
    """
    pass

  @property
  def eval_dataset(self):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
402
    """The current evaluation dataset."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
403
404
405
406
    return self._eval_dataset

  @eval_dataset.setter
  def eval_dataset(self, eval_dataset):
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
407
408
409
    """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
410
411

    Args:
Dan Holtmann-Rice's avatar
Dan Holtmann-Rice committed
412
413
      eval_dataset: A `tf.nest`-compatible structure of `tf.data.Dataset` or
        `DistributedDataset`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
414
415
    """
    self._eval_dataset = eval_dataset