keras_common.py 9.21 KB
Newer Older
1
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2
3
4
5
6
7
8
9
10
11
12
13
14
#
# 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
"""Common util functions and classes used by both keras cifar and imagenet."""
16
17
18
19
20
21
22

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

import time

23
24
import numpy as np

Toby Boyd's avatar
Toby Boyd committed
25
26
27
# pylint: disable=g-bad-import-order
from absl import flags
import tensorflow as tf
28
29
from tensorflow.python.keras.optimizer_v2 import (gradient_descent as
                                                  gradient_descent_v2)
30

Shining Sun's avatar
Shining Sun committed
31
FLAGS = flags.FLAGS
Shining Sun's avatar
Shining Sun committed
32
BASE_LEARNING_RATE = 0.1  # This matches Jing's version.
33
34
TRAIN_TOP_1 = 'training_accuracy_top_1'

Shining Sun's avatar
Shining Sun committed
35

36
37
38
39
40
41
42
43
class BatchTimestamp(object):
  """A structure to store batch time stamp."""

  def __init__(self, batch_index, timestamp):
    self.batch_index = batch_index
    self.timestamp = timestamp


44
45
46
class TimeHistory(tf.keras.callbacks.Callback):
  """Callback for Keras models."""

47
  def __init__(self, batch_size, log_steps):
48
    """Callback for logging performance (# image/second).
49
50
51
52
53

    Args:
      batch_size: Total batch size.

    """
54
    self.batch_size = batch_size
55
    super(TimeHistory, self).__init__()
56
57
    self.log_steps = log_steps

58
59
    # Logs start of step 0 then end of each step based on log_steps interval.
    self.timestamp_log = []
60
61
62
63

  def on_train_begin(self, logs=None):
    self.record_batch = True

64
65
66
  def on_train_end(self, logs=None):
    self.train_finish_time = time.time()

67
68
  def on_batch_begin(self, batch, logs=None):
    if self.record_batch:
69
70
      timestamp = time.time()
      self.start_time = timestamp
71
      self.record_batch = False
72
73
      if batch == 0:
        self.timestamp_log.append(BatchTimestamp(batch, timestamp))
74
75

  def on_batch_end(self, batch, logs=None):
Shining Sun's avatar
Shining Sun committed
76
    if batch % self.log_steps == 0:
77
78
      timestamp = time.time()
      elapsed_time = timestamp - self.start_time
79
      examples_per_second = (self.batch_size * self.log_steps) / elapsed_time
80
      if batch != 0:
81
82
        self.record_batch = True
        self.timestamp_log.append(BatchTimestamp(batch, timestamp))
83
84
85
86
        tf.compat.v1.logging.info(
            "BenchmarkMetric: {'num_batches':%d, 'time_taken': %f,"
            "'images_per_second': %f}" %
            (batch, elapsed_time, examples_per_second))
87

88

89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
class LearningRateBatchScheduler(tf.keras.callbacks.Callback):
  """Callback to update learning rate on every batch (not epoch boundaries).

  N.B. Only support Keras optimizers, not TF optimizers.

  Args:
      schedule: a function that takes an epoch index and a batch index as input
          (both integer, indexed from 0) and returns a new learning rate as
          output (float).
  """

  def __init__(self, schedule, batch_size, num_images):
    super(LearningRateBatchScheduler, self).__init__()
    self.schedule = schedule
    self.batches_per_epoch = num_images / batch_size
    self.batch_size = batch_size
    self.epochs = -1
    self.prev_lr = -1

  def on_epoch_begin(self, epoch, logs=None):
109
110
    if not hasattr(self.model.optimizer, 'learning_rate'):
      raise ValueError('Optimizer must have a "learning_rate" attribute.')
111
112
113
    self.epochs += 1

  def on_batch_begin(self, batch, logs=None):
114
    """Executes before step begins."""
115
116
117
118
    lr = self.schedule(self.epochs,
                       batch,
                       self.batches_per_epoch,
                       self.batch_size)
119
120
121
    if not isinstance(lr, (float, np.float32, np.float64)):
      raise ValueError('The output of the "schedule" function should be float.')
    if lr != self.prev_lr:
Shining Sun's avatar
Shining Sun committed
122
      self.model.optimizer.learning_rate = lr  # lr should be a float here
123
      self.prev_lr = lr
124
125
126
      tf.compat.v1.logging.debug(
          'Epoch %05d Batch %05d: LearningRateBatchScheduler '
          'change learning rate to %s.', self.epochs, batch, lr)
127

128

Shining Sun's avatar
Shining Sun committed
129
def get_optimizer():
130
131
  """Returns optimizer to use."""
  # The learning_rate is overwritten at the beginning of each step by callback.
Shining Sun's avatar
Shining Sun committed
132
  return gradient_descent_v2.SGD(learning_rate=0.1, momentum=0.9)
133
134


135
def get_callbacks(learning_rate_schedule_fn, num_images):
136
  """Returns common callbacks."""
137
  time_callback = TimeHistory(FLAGS.batch_size, FLAGS.log_steps)
138
139

  tensorboard_callback = tf.keras.callbacks.TensorBoard(
140
      log_dir=FLAGS.model_dir)
141

Shining Sun's avatar
Shining Sun committed
142
  lr_callback = LearningRateBatchScheduler(
143
144
145
      learning_rate_schedule_fn,
      batch_size=FLAGS.batch_size,
      num_images=num_images)
146
147
148

  return time_callback, tensorboard_callback, lr_callback

Shining Sun's avatar
Shining Sun committed
149

150
def build_stats(history, eval_output, time_callback):
151
152
153
154
155
156
157
  """Normalizes and returns dictionary of stats.

  Args:
    history: Results of the training step. Supports both categorical_accuracy
      and sparse_categorical_accuracy.
    eval_output: Output of the eval step. Assumes first value is eval_loss and
      second value is accuracy_top_1.
158
    time_callback: Time tracking callback likely used during keras.fit.
159
160
161
162
163
164
165
166

  Returns:
    Dictionary of normalized results.
  """
  stats = {}
  if eval_output:
    stats['accuracy_top_1'] = eval_output[1].item()
    stats['eval_loss'] = eval_output[0].item()
167

168
169
170
171
172
173
174
175
176
177
  if history and history.history:
    train_hist = history.history
    # Gets final loss from training.
    stats['loss'] = train_hist['loss'][-1].item()
    # Gets top_1 training accuracy.
    if 'categorical_accuracy' in train_hist:
      stats[TRAIN_TOP_1] = train_hist['categorical_accuracy'][-1].item()
    elif 'sparse_categorical_accuracy' in train_hist:
      stats[TRAIN_TOP_1] = train_hist['sparse_categorical_accuracy'][-1].item()

178
  if time_callback:
179
180
    timestamp_log = time_callback.timestamp_log
    stats['step_timestamp_log'] = timestamp_log
181
    stats['train_finish_time'] = time_callback.train_finish_time
182
183
184
185
186
    if len(timestamp_log) > 1:
      stats['avg_exp_per_second'] = (
          time_callback.batch_size * time_callback.log_steps *
          (len(time_callback.timestamp_log)-1) /
          (timestamp_log[-1].timestamp - timestamp_log[0].timestamp))
187

188
189
190
  return stats


Shining Sun's avatar
Shining Sun committed
191
192
def define_keras_flags():
  flags.DEFINE_boolean(name='enable_eager', default=False, help='Enable eager?')
193
  flags.DEFINE_boolean(name='skip_eval', default=False, help='Skip evaluation?')
Shining Sun's avatar
Shining Sun committed
194
  flags.DEFINE_integer(
195
196
      name='train_steps', default=None,
      help='The number of steps to run for training. If it is larger than '
Shining Sun's avatar
Shining Sun committed
197
      '# batches per epoch, then use # batches per epoch. When this flag is '
198
      'set, only one epoch is going to run for training.')
199
200
201
202
203
  flags.DEFINE_integer(
      name='log_steps', default=100,
      help='For every log_steps, we log the timing information such as '
      'examples per second. Besides, for every log_steps, we store the '
      'timestamp of a batch end.')
204

Shining Sun's avatar
Shining Sun committed
205
206
207
208
209
210
211
212

def get_synth_input_fn(height, width, num_channels, num_classes,
                       dtype=tf.float32):
  """Returns an input function that returns a dataset with random data.

  This input_fn returns a data set that iterates over a set of random data and
  bypasses all preprocessing, e.g. jpeg decode and copy. The host to device
  copy is still included. This used to find the upper throughput bound when
Shining Sun's avatar
Shining Sun committed
213
  tuning the full input pipeline.
Shining Sun's avatar
Shining Sun committed
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230

  Args:
    height: Integer height that will be used to create a fake image tensor.
    width: Integer width that will be used to create a fake image tensor.
    num_channels: Integer depth that will be used to create a fake image tensor.
    num_classes: Number of classes that should be represented in the fake labels
      tensor
    dtype: Data type for features/images.

  Returns:
    An input_fn that can be used in place of a real one to return a dataset
    that can be used for iteration.
  """
  # pylint: disable=unused-argument
  def input_fn(is_training, data_dir, batch_size, *args, **kwargs):
    """Returns dataset filled with random data."""
    # Synthetic input should be within [0, 255].
231
232
233
234
235
236
237
238
239
240
241
    inputs = tf.random.truncated_normal([height, width, num_channels],
                                        dtype=dtype,
                                        mean=127,
                                        stddev=60,
                                        name='synthetic_inputs')

    labels = tf.random.uniform([1],
                               minval=0,
                               maxval=num_classes - 1,
                               dtype=tf.int32,
                               name='synthetic_labels')
Shining Sun's avatar
Shining Sun committed
242
    data = tf.data.Dataset.from_tensors((inputs, labels)).repeat()
243
    data = data.batch(batch_size)
244
    data = data.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
Shining Sun's avatar
Shining Sun committed
245
246
247
    return data

  return input_fn
Shining Sun's avatar
Shining Sun committed
248
249
250
251
252
253


def get_strategy_scope(strategy):
  if strategy:
    strategy_scope = strategy.scope()
  else:
Shining Sun's avatar
Shining Sun committed
254
    strategy_scope = DummyContextManager()
Shining Sun's avatar
Shining Sun committed
255
256
257
258
259

  return strategy_scope


class DummyContextManager(object):
Shining Sun's avatar
Shining Sun committed
260

Shining Sun's avatar
Shining Sun committed
261
262
263
264
265
  def __enter__(self):
    pass

  def __exit__(self, *args):
    pass