model_training_utils.py 24.9 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
# Copyright 2019 The TensorFlow 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.
# ==============================================================================
15
"""A light weight utilities to train NLP models."""
16
17
18
19
20

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

21
import json
22
import os
23
import tempfile
24
25

from absl import logging
Hongkun Yu's avatar
Hongkun Yu committed
26
import tensorflow as tf
27
from tensorflow.python.util import deprecation
28
from official.common import distribute_utils
Zongwei Zhou's avatar
Zongwei Zhou committed
29
from official.staging.training import grad_utils
30

31
32
_SUMMARY_TXT = 'training_summary.txt'
_MIN_SUMMARY_STEPS = 10
33

34

35
36
37
38
39
40
41
42
43
def _should_export_checkpoint(strategy):
  return (not strategy) or strategy.extended.should_checkpoint


def _should_export_summary(strategy):
  return (not strategy) or strategy.extended.should_save_summary


def _save_checkpoint(strategy, checkpoint, model_dir, checkpoint_prefix):
44
45
  """Saves model to with provided checkpoint prefix."""

46
47
48
49
50
51
52
53
54
55
56
57
  if _should_export_checkpoint(strategy):
    checkpoint_path = os.path.join(model_dir, checkpoint_prefix)
    saved_path = checkpoint.save(checkpoint_path)
    logging.info('Saving model as TF checkpoint: %s', saved_path)
  else:
    # In multi worker training we need every worker to save checkpoint, because
    # variables can trigger synchronization on read and synchronization needs
    # all workers to participate. To avoid workers overriding each other we save
    # to a temporary directory on non-chief workers.
    tmp_dir = tempfile.mkdtemp()
    checkpoint.save(os.path.join(tmp_dir, 'ckpt'))
    tf.io.gfile.rmtree(tmp_dir)
58
59
60
  return


61
62
63
64
65
def _get_input_iterator(input_fn, strategy):
  """Returns distributed dataset iterator."""
  # When training with TPU pods, datasets needs to be cloned across
  # workers. Since Dataset instance cannot be cloned in eager mode, we instead
  # pass callable that returns a dataset.
Hongkun Yu's avatar
Hongkun Yu committed
66
67
  if not callable(input_fn):
    raise ValueError('`input_fn` should be a closure that returns a dataset.')
Chenkai Kuang's avatar
Chenkai Kuang committed
68
  iterator = iter(strategy.distribute_datasets_from_function(input_fn))
69
70
71
  return iterator


72
73
74
75
76
def _float_metric_value(metric):
  """Gets the value of a float-value keras metric."""
  return metric.result().numpy().astype(float)


77
78
79
80
81
82
83
def clip_by_global_norm_callback(grads_and_vars):
  """Performs gradient clipping."""
  grads, variables = zip(*grads_and_vars)
  (clipped_grads, _) = tf.clip_by_global_norm(grads, clip_norm=1.0)
  return zip(clipped_grads, variables)


84
def steps_to_run(current_step, steps_per_epoch, steps_per_loop):
85
  """Calculates steps to run on device."""
86
87
88
  if steps_per_loop <= 0:
    raise ValueError('steps_per_loop should be positive integer.')
  if steps_per_loop == 1:
89
90
91
92
93
94
95
96
    return steps_per_loop
  remainder_in_epoch = current_step % steps_per_epoch
  if remainder_in_epoch != 0:
    return min(steps_per_epoch - remainder_in_epoch, steps_per_loop)
  else:
    return steps_per_loop


97
def write_txt_summary(training_summary, summary_dir):
98
  """Writes a summary text file to record stats."""
Chen Chen's avatar
Chen Chen committed
99
100
  if not tf.io.gfile.exists(summary_dir):
    tf.io.gfile.mkdir(summary_dir)
101
  summary_path = os.path.join(summary_dir, _SUMMARY_TXT)
102
103
104
105
106
  with tf.io.gfile.GFile(summary_path, 'wb') as f:
    logging.info('Training Summary: \n%s', str(training_summary))
    f.write(json.dumps(training_summary, indent=4))


107
@deprecation.deprecated(
108
109
110
    None, 'This function is deprecated and we do not expect adding new '
    'functionalities. Please do not have your code depending '
    'on this library.')
111
112
113
114
115
116
117
def run_customized_training_loop(
    # pylint: disable=invalid-name
    _sentinel=None,
    # pylint: enable=invalid-name
    strategy=None,
    model_fn=None,
    loss_fn=None,
118
    scale_loss=True,
119
120
121
    model_dir=None,
    train_input_fn=None,
    steps_per_epoch=None,
Tianqi Liu's avatar
Tianqi Liu committed
122
    num_eval_per_epoch=1,
123
    steps_per_loop=None,
124
125
126
127
128
    epochs=1,
    eval_input_fn=None,
    eval_steps=None,
    metric_fn=None,
    init_checkpoint=None,
129
    custom_callbacks=None,
Chen Chen's avatar
Chen Chen committed
130
    run_eagerly=False,
Zongwei Zhou's avatar
Zongwei Zhou committed
131
132
133
    sub_model_export_name=None,
    explicit_allreduce=False,
    pre_allreduce_callbacks=None,
Chen Chen's avatar
Chen Chen committed
134
    post_allreduce_callbacks=None,
Zongwei Zhou's avatar
Zongwei Zhou committed
135
136
    train_summary_interval=0,
    allreduce_bytes_per_pack=0):
137
138
139
140
141
142
143
144
145
146
147
148
  """Run BERT pretrain model training using low-level API.

  Arguments:
      _sentinel: Used to prevent positional parameters. Internal, do not use.
      strategy: Distribution strategy on which to run low level training loop.
      model_fn: Function that returns a tuple (model, sub_model). Caller of this
        function should add optimizer to the `model` via calling
        `model.compile()` API or manually setting `model.optimizer` attribute.
        Second element of the returned tuple(sub_model) is an optional sub model
        to be used for initial checkpoint -- if provided.
      loss_fn: Function with signature func(labels, logits) and returns a loss
        tensor.
149
150
      scale_loss: Whether to divide the raw loss by number of replicas before
        gradients calculation.
151
152
153
      model_dir: Model directory used during training for restoring/saving model
        weights.
      train_input_fn: Function that returns a tf.data.Dataset used for training.
154
155
156
      steps_per_epoch: Number of steps to run per epoch. At the end of each
        epoch, model checkpoint will be saved and evaluation will be conducted
        if evaluation dataset is provided.
Tianqi Liu's avatar
Tianqi Liu committed
157
      num_eval_per_epoch: Number of evaluations per epoch.
158
159
160
      steps_per_loop: Number of steps per graph-mode loop. In order to reduce
        communication in eager context, training logs are printed every
        steps_per_loop.
161
162
163
164
165
      epochs: Number of epochs to train.
      eval_input_fn: Function that returns evaluation dataset. If none,
        evaluation is skipped.
      eval_steps: Number of steps to run evaluation. Required if `eval_input_fn`
        is not none.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
166
167
168
      metric_fn: A metrics function that returns either a Keras Metric object or
        a list of Keras Metric objects to record evaluation result using
        evaluation dataset or with training dataset after every epoch.
169
170
      init_checkpoint: Optional checkpoint to load to `sub_model` returned by
        `model_fn`.
171
      custom_callbacks: A list of Keras Callbacks objects to run during
172
173
        training. More specifically, `on_train_begin(), on_train_end(),
        on_batch_begin()`, `on_batch_end()`, `on_epoch_begin()`,
Hongkun Yu's avatar
Hongkun Yu committed
174
175
        `on_epoch_end()` methods are invoked during training. Note that some
        metrics may be missing from `logs`.
176
177
      run_eagerly: Whether to run model training in pure eager execution. This
        should be disable for TPUStrategy.
Chen Chen's avatar
Chen Chen committed
178
179
180
      sub_model_export_name: If not None, will export `sub_model` returned by
        `model_fn` into checkpoint files. The name of intermediate checkpoint
        file is {sub_model_export_name}_step_{step}.ckpt and the last
Tianqi Liu's avatar
Tianqi Liu committed
181
182
        checkpint's name is {sub_model_export_name}.ckpt; if None, `sub_model`
        will not be exported as checkpoint.
Zongwei Zhou's avatar
Zongwei Zhou committed
183
184
185
186
187
188
189
190
191
      explicit_allreduce: Whether to explicitly perform gradient allreduce,
        instead of relying on implicit allreduce in optimizer.apply_gradients().
        default is False. For now, if training using FP16 mixed precision,
        explicit allreduce will aggregate gradients in FP16 format. For TPU and
        GPU training using FP32, explicit allreduce will aggregate gradients in
        FP32 format.
      pre_allreduce_callbacks: A list of callback functions that takes gradients
        and model variables pairs as input, manipulate them, and returns a new
        gradients and model variables paris. The callback functions will be
Tianqi Liu's avatar
Tianqi Liu committed
192
193
194
195
        invoked in the list order and before gradients are allreduced. With
        mixed precision training, the pre_allreduce_allbacks will be applied on
        scaled_gradients. Default is no callbacks. Only used when
        explicit_allreduce=True.
Zongwei Zhou's avatar
Zongwei Zhou committed
196
197
198
199
200
201
      post_allreduce_callbacks: A list of callback functions that takes
        gradients and model variables pairs as input, manipulate them, and
        returns a new gradients and model variables paris. The callback
        functions will be invoked in the list order and right before gradients
        are applied to variables for updates. Default is no callbacks. Only used
        when explicit_allreduce=True.
Chen Chen's avatar
Chen Chen committed
202
203
      train_summary_interval: Step interval for training summaries. If the value
        is a negative number, then training summaries are not enabled.
Zongwei Zhou's avatar
Zongwei Zhou committed
204
205
206
207
208
      allreduce_bytes_per_pack: A non-negative integer. Breaks collective
        operations into packs of certain size. If it's zero, all gradients are
        in one pack. Breaking gradient into packs could enable overlap between
        allreduce and backprop computation. This flag only takes effect when
        explicit_allreduce is set to True.'
209
210
211
212
213
214
215
216

  Returns:
      Trained model.

  Raises:
      ValueError: (1) When model returned by `model_fn` does not have optimizer
        attribute or when required parameters are set to none. (2) eval args are
        not specified correctly. (3) metric_fn must be a callable if specified.
Chen Chen's avatar
Chen Chen committed
217
218
        (4) sub_model_checkpoint_name is specified, but `sub_model` returned
        by `model_fn` is None.
219
220
221
222
223
224
225
226
227
  """

  if _sentinel is not None:
    raise ValueError('only call `run_customized_training_loop()` '
                     'with named arguments.')

  required_arguments = [
      strategy, model_fn, loss_fn, model_dir, steps_per_epoch, train_input_fn
  ]
Tianqi Liu's avatar
Tianqi Liu committed
228
229

  steps_between_evals = int(steps_per_epoch / num_eval_per_epoch)
230
231
  if [arg for arg in required_arguments if arg is None]:
    raise ValueError('`strategy`, `model_fn`, `loss_fn`, `model_dir`, '
232
                     '`steps_per_epoch` and `train_input_fn` are required '
233
                     'parameters.')
234
235
236
237
  if not steps_per_loop:
    if tf.config.list_logical_devices('TPU'):
      # One can't fully utilize a TPU with steps_per_loop=1, so in this case
      # default users to a more useful value.
Tianqi Liu's avatar
Tianqi Liu committed
238
      steps_per_loop = min(1000, steps_between_evals)
239
240
241
242
    else:
      steps_per_loop = 1
    logging.info('steps_per_loop not specified. Using steps_per_loop=%d',
                 steps_per_loop)
Tianqi Liu's avatar
Tianqi Liu committed
243
  if steps_per_loop > steps_between_evals:
244
    logging.warning(
245
        'steps_per_loop: %d is specified to be greater than '
Tianqi Liu's avatar
Tianqi Liu committed
246
247
248
        ' steps_between_evals: %d, we will use steps_between_evals as'
        ' steps_per_loop.', steps_per_loop, steps_between_evals)
    steps_per_loop = steps_between_evals
249
250
  assert tf.executing_eagerly()

251
252
253
  if run_eagerly:
    if isinstance(strategy, tf.distribute.experimental.TPUStrategy):
      raise ValueError(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
254
          'TPUStrategy should not run eagerly as it heavily relies on graph'
255
256
          ' optimization for the distributed system.')

257
  if eval_input_fn and eval_steps is None:
258
    raise ValueError(
259
        '`eval_step` is required when `eval_input_fn ` is not none.')
260
261
262
263
  if metric_fn and not callable(metric_fn):
    raise ValueError(
        'if `metric_fn` is specified, metric_fn must be a callable.')

264
  total_training_steps = steps_per_epoch * epochs
265
  train_iterator = _get_input_iterator(train_input_fn, strategy)
Tianqi Liu's avatar
Tianqi Liu committed
266
  eval_loss_metric = tf.keras.metrics.Mean('training_loss', dtype=tf.float32)
267

268
  with distribute_utils.get_strategy_scope(strategy):
269
270
271
272
273
274
    # To correctly place the model weights on accelerators,
    # model and optimizer should be created in scope.
    model, sub_model = model_fn()
    if not hasattr(model, 'optimizer'):
      raise ValueError('User should set optimizer attribute to model '
                       'inside `model_fn`.')
Chen Chen's avatar
Chen Chen committed
275
276
277
278
    if sub_model_export_name and sub_model is None:
      raise ValueError('sub_model_export_name is specified as %s, but '
                       'sub_model is None.' % sub_model_export_name)

279
280
281
    callback_list = tf.keras.callbacks.CallbackList(
        callbacks=custom_callbacks, model=model)

282
283
284
285
286
287
288
    optimizer = model.optimizer

    if init_checkpoint:
      logging.info(
          'Checkpoint file %s found and restoring from '
          'initial checkpoint for core model.', init_checkpoint)
      checkpoint = tf.train.Checkpoint(model=sub_model)
Jing Li's avatar
Jing Li committed
289
      checkpoint.restore(init_checkpoint).assert_existing_objects_matched()
290
291
      logging.info('Loading from checkpoint file completed')

Tianqi Liu's avatar
Tianqi Liu committed
292
    train_loss_metric = tf.keras.metrics.Mean('training_loss', dtype=tf.float32)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
293
294
295
    eval_metrics = metric_fn() if metric_fn else []
    if not isinstance(eval_metrics, list):
      eval_metrics = [eval_metrics]
296
297
298
299
300
301
302
303
    # If evaluation is required, make a copy of metric as it will be used by
    # both train and evaluation.
    train_metrics = [
        metric.__class__.from_config(metric.get_config())
        for metric in eval_metrics
    ]

    # Create summary writers
304
305
306
307
308
309
310
    if _should_export_summary(strategy):
      summary_dir = os.path.join(model_dir, 'summaries')
    else:
      # In multi worker training we need every worker to write summary, because
      # variables can trigger synchronization on read and synchronization needs
      # all workers to participate.
      summary_dir = tempfile.mkdtemp()
311
    eval_summary_writer = tf.summary.create_file_writer(
312
        os.path.join(summary_dir, 'eval'))
Chen Chen's avatar
Chen Chen committed
313
314
    last_summary_step = 0
    if steps_per_loop >= _MIN_SUMMARY_STEPS and train_summary_interval >= 0:
315
316
317
      # Only writes summary when the stats are collected sufficiently over
      # enough steps.
      train_summary_writer = tf.summary.create_file_writer(
318
          os.path.join(summary_dir, 'train'))
319
    else:
Chen Chen's avatar
Chen Chen committed
320
      train_summary_writer = tf.summary.create_noop_writer()
321
322
323
324
325
326
327
328
329
330
331

    # Collects training variables.
    training_vars = model.trainable_variables

    def _replicated_step(inputs):
      """Replicated training step."""

      inputs, labels = inputs
      with tf.GradientTape() as tape:
        model_outputs = model(inputs, training=True)
        loss = loss_fn(labels, model_outputs)
332
333
334
335
336
337
        # Raw loss is used for reporting in metrics/logs.
        raw_loss = loss
        if scale_loss:
          # Scales down the loss for gradients to be invariant from replicas.
          loss = loss / strategy.num_replicas_in_sync

Zongwei Zhou's avatar
Zongwei Zhou committed
338
339
340
341
      if explicit_allreduce:
        grad_utils.minimize_using_explicit_allreduce(tape, optimizer, loss,
                                                     training_vars,
                                                     pre_allreduce_callbacks,
Zongwei Zhou's avatar
Zongwei Zhou committed
342
343
                                                     post_allreduce_callbacks,
                                                     allreduce_bytes_per_pack)
344
      else:
Zongwei Zhou's avatar
Zongwei Zhou committed
345
346
347
348
349
350
351
352
353
        if isinstance(optimizer,
                      tf.keras.mixed_precision.experimental.LossScaleOptimizer):
          with tape:
            scaled_loss = optimizer.get_scaled_loss(loss)
          scaled_grads = tape.gradient(scaled_loss, training_vars)
          grads = optimizer.get_unscaled_gradients(scaled_grads)
        else:
          grads = tape.gradient(loss, training_vars)
        optimizer.apply_gradients(zip(grads, training_vars))
354
      # For reporting, the metric takes the mean of losses.
355
      train_loss_metric.update_state(raw_loss)
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
      for metric in train_metrics:
        metric.update_state(labels, model_outputs)

    @tf.function
    def train_steps(iterator, steps):
      """Performs distributed training steps in a loop.

      Args:
        iterator: the distributed iterator of training datasets.
        steps: an tf.int32 integer tensor to specify number of steps to run
          inside host training loop.

      Raises:
        ValueError: Any of the arguments or tensor shapes are invalid.
      """
      if not isinstance(steps, tf.Tensor):
        raise ValueError('steps should be an Tensor. Python object may cause '
                         'retracing.')

      for _ in tf.range(steps):
Ken Franko's avatar
Ken Franko committed
376
        strategy.run(_replicated_step, args=(next(iterator),))
377

378
379
    def train_single_step(iterator):
      """Performs a distributed training step.
380

381
382
      Args:
        iterator: the distributed iterator of training datasets.
383

384
385
386
      Raises:
        ValueError: Any of the arguments or tensor shapes are invalid.
      """
Ken Franko's avatar
Ken Franko committed
387
      strategy.run(_replicated_step, args=(next(iterator),))
388

389
390
    def test_step(iterator):
      """Calculates evaluation metrics on distributed devices."""
391

392
393
      def _test_step_fn(inputs):
        """Replicated accuracy calculation."""
394

395
396
397
398
        inputs, labels = inputs
        model_outputs = model(inputs, training=False)
        for metric in eval_metrics:
          metric.update_state(labels, model_outputs)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
399
        return model_outputs, labels
400

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
401
402
403
404
405
406
      outputs, labels = strategy.run(_test_step_fn, args=(next(iterator),))
      outputs = tf.nest.map_structure(strategy.experimental_local_results,
                                      outputs)
      labels = tf.nest.map_structure(strategy.experimental_local_results,
                                     labels)
      return outputs, labels
407
408
409
410
411
412

    if not run_eagerly:
      train_single_step = tf.function(train_single_step)
      test_step = tf.function(test_step)

    def _run_evaluation(current_training_step, test_iterator):
413
414
415
416
417
418
419
420
421
      """Runs validation steps and aggregate metrics.

      Args:
        current_training_step: tf.int32 tensor containing the current step.
        test_iterator: distributed iterator of test datasets.

      Returns:
        A dict of metic names and values.
      """
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
422
423
424
425
426
427
      # The last batch of the evaluation is often smaller than previous ones.
      # Moreover, in some distributed pieces it might even be empty. Therefore,
      # different from the way training_loss is calculated, it is needed to
      # gather all the logits and labels here to calculate the evaluation loss
      # outside.
      loss_list, loss_weights = list(), list()
428
      for _ in range(eval_steps):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
429
430
431
432
433
434
435
436
437
438
439
440
        outputs, labels = test_step(test_iterator)
        for cur_logits, cur_labels in zip(outputs, labels):
          # This is to handle cases when cur_labels is not a single tensor,
          # but a dict of tensors.
          cur_weight = tf.shape(tf.nest.flatten(cur_labels)[0])[0]
          if cur_weight != 0:
            loss_list.append(loss_fn(cur_labels, cur_logits).numpy())
            loss_weights.append(cur_weight)
      # The sample_weights are the actual number of examples in each batch,
      # a summation of numbers of examples in each replica if using
      # distributed training.
      eval_loss_metric.update_state(loss_list, sample_weight=loss_weights)
441

442
      logs = {}
443
      with eval_summary_writer.as_default():
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
444
        for metric in [eval_loss_metric] + eval_metrics + model.metrics:
445
          metric_value = _float_metric_value(metric)
446
          logs[metric.name] = metric_value
447
448
449
450
451
452
          logging.info('Step: [%d] Validation %s = %f', current_training_step,
                       metric.name, metric_value)
          tf.summary.scalar(
              metric.name, metric_value, step=current_training_step)
        eval_summary_writer.flush()

453
      return logs
454
455

    # Training loop starts here.
Le Hou's avatar
Le Hou committed
456
457
    checkpoint = tf.train.Checkpoint(
        model=model, optimizer=optimizer, global_step=optimizer.iterations)
Chen Chen's avatar
Chen Chen committed
458
    sub_model_checkpoint = tf.train.Checkpoint(
Le Hou's avatar
Le Hou committed
459
460
        model=sub_model,
        global_step=optimizer.iterations) if sub_model_export_name else None
Chen Chen's avatar
Chen Chen committed
461

462
463
    latest_checkpoint_file = tf.train.latest_checkpoint(model_dir)
    if latest_checkpoint_file:
Tianqi Liu's avatar
Tianqi Liu committed
464
465
      logging.info('Checkpoint file %s found and restoring from '
                   'checkpoint', latest_checkpoint_file)
466
467
468
469
470
471
      checkpoint.restore(latest_checkpoint_file)
      logging.info('Loading from checkpoint file completed')

    current_step = optimizer.iterations.numpy()
    checkpoint_name = 'ctl_step_{step}.ckpt'

Tianqi Liu's avatar
Tianqi Liu committed
472
    logs = {}
473
474
    callback_list.on_train_begin()
    while current_step < total_training_steps and not model.stop_training:
475
      if current_step % steps_per_epoch == 0:
Hongkun Yu's avatar
Hongkun Yu committed
476
        callback_list.on_epoch_begin(int(current_step / steps_per_epoch) + 1)
477

478
479
480
481
482
483
      # Training loss/metric are taking average over steps inside micro
      # training loop. We reset the their values before each round.
      train_loss_metric.reset_states()
      for metric in train_metrics + model.metrics:
        metric.reset_states()

484
      callback_list.on_batch_begin(current_step)
485
      # Runs several steps in the host while loop.
Tianqi Liu's avatar
Tianqi Liu committed
486
      steps = steps_to_run(current_step, steps_between_evals, steps_per_loop)
487

488
      if tf.config.list_physical_devices('GPU'):
489
490
        # TODO(zongweiz): merge with train_steps once tf.while_loop
        # GPU performance bugs are fixed.
491
492
        for _ in range(steps):
          train_single_step(train_iterator)
493
494
      else:
        # Converts steps to a Tensor to avoid tf.function retracing.
Tianqi Liu's avatar
Tianqi Liu committed
495
        train_steps(train_iterator, tf.convert_to_tensor(steps, dtype=tf.int32))
496
      train_loss = _float_metric_value(train_loss_metric)
497
498
499
500
501
502
      current_step += steps

      # Updates training logging.
      training_status = 'Train Step: %d/%d  / loss = %s' % (
          current_step, total_training_steps, train_loss)

Chen Chen's avatar
Chen Chen committed
503
504
505
506
507
508
509
      if current_step >= last_summary_step + train_summary_interval:
        summary_writer = train_summary_writer
        last_summary_step = current_step
      else:
        summary_writer = tf.summary.create_noop_writer()

      with summary_writer.as_default():
510
511
512
513
514
        if callable(optimizer.learning_rate):
          tf.summary.scalar(
              'learning_rate',
              optimizer.learning_rate(current_step),
              step=current_step)
Tianqi Liu's avatar
Tianqi Liu committed
515
        tf.summary.scalar(train_loss_metric.name, train_loss, step=current_step)
Chen Chen's avatar
Chen Chen committed
516
517
518
519
520
        for metric in train_metrics + model.metrics:
          metric_value = _float_metric_value(metric)
          training_status += '  %s = %f' % (metric.name, metric_value)
          tf.summary.scalar(metric.name, metric_value, step=current_step)
        summary_writer.flush()
521
522
      logging.info(training_status)

Tianqi Liu's avatar
Tianqi Liu committed
523
524
525
526
527
      # If no need for evaluation, we only call on_batch_end with train_loss,
      # this is to ensure we get granular global_step/sec on Tensorboard.
      if current_step % steps_between_evals:
        callback_list.on_batch_end(current_step - 1, {'loss': train_loss})
      else:
528
529
530
531
532
533
534
535
536
        # Save a submodel with the step in the file name after each epoch.
        if sub_model_export_name:
          _save_checkpoint(
              strategy, sub_model_checkpoint, model_dir,
              '%s_step_%d.ckpt' % (sub_model_export_name, current_step))

        # Save model checkpoints and run validation steps after each epoch
        # (with the exception of the final epoch which is handled after the
        # training loop).
537
        if current_step < total_training_steps:
538
          _save_checkpoint(strategy, checkpoint, model_dir,
539
                           checkpoint_name.format(step=current_step))
540
541
          if eval_input_fn:
            # Re-initialize evaluation metric.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
542
            eval_loss_metric.reset_states()
543
544
            for metric in eval_metrics + model.metrics:
              metric.reset_states()
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
545
546
547
548

            logging.info('Running evaluation after step: %s.', current_step)
            logs = _run_evaluation(current_step,
                                   _get_input_iterator(eval_input_fn, strategy))
Tianqi Liu's avatar
Tianqi Liu committed
549
550
551
552
        # We add train_loss here rather than call on_batch_end twice to make
        # sure that no duplicated values are generated.
        logs['loss'] = train_loss
        callback_list.on_batch_end(current_step - 1, logs)
553

Tianqi Liu's avatar
Tianqi Liu committed
554
555
556
557
      # Calls on_epoch_end after each real epoch ends to prevent mis-calculation
      # of training steps.
      if current_step % steps_per_epoch == 0:
        callback_list.on_epoch_end(int(current_step / steps_per_epoch), logs)
558

Chen Chen's avatar
Chen Chen committed
559
    if sub_model_export_name:
560
      _save_checkpoint(strategy, sub_model_checkpoint, model_dir,
Chen Chen's avatar
Chen Chen committed
561
                       '%s.ckpt' % sub_model_export_name)
562

563
564
    _save_checkpoint(strategy, checkpoint, model_dir,
                     checkpoint_name.format(step=current_step))
565
    if eval_input_fn:
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
566
567
568
569
570
      # Re-initialize evaluation metric.
      eval_loss_metric.reset_states()
      for metric in eval_metrics + model.metrics:
        metric.reset_states()

571
      logging.info('Running final evaluation after training is complete.')
572
573
574
      logs = _run_evaluation(current_step,
                             _get_input_iterator(eval_input_fn, strategy))
    callback_list.on_epoch_end(int(current_step / steps_per_epoch), logs)
575
576
577
578
    training_summary = {
        'total_training_steps': total_training_steps,
        'train_loss': _float_metric_value(train_loss_metric),
    }
579
580
    for metric in model.metrics:
      training_summary[metric.name] = _float_metric_value(metric)
581
582
583
584
    if eval_metrics:
      training_summary['last_train_metrics'] = _float_metric_value(
          train_metrics[0])
      training_summary['eval_metrics'] = _float_metric_value(eval_metrics[0])
585

586
    write_txt_summary(training_summary, summary_dir)
587

588
589
590
    if not _should_export_summary(strategy):
      tf.io.gfile.rmtree(summary_dir)

591
592
    callback_list.on_train_end()

593
    return model