keras_common.py 7.36 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
23

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

import time

from absl import flags
Toby Boyd's avatar
Toby Boyd committed
24
import numpy as np  # pylint: disable=g-bad-import-order
25
26
27
28
29
import tensorflow as tf  # pylint: disable=g-bad-import-order

from tensorflow.python.keras.optimizer_v2 import gradient_descent as gradient_descent_v2


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

Shining Sun's avatar
Shining Sun committed
34

35
36
37
38
class TimeHistory(tf.keras.callbacks.Callback):
  """Callback for Keras models."""

  def __init__(self, batch_size):
39
    """Callback for logging performance (# image/second).
40
41
42
43
44
45
46

    Args:
      batch_size: Total batch size.

    """
    self._batch_size = batch_size
    super(TimeHistory, self).__init__()
Shining Sun's avatar
Shining Sun committed
47
    self.log_steps = 100
48
49
50
51
52
53

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

  def on_batch_begin(self, batch, logs=None):
    if self.record_batch:
54
      self.start_time = time.time()
55
56
57
      self.record_batch = False

  def on_batch_end(self, batch, logs=None):
Shining Sun's avatar
Shining Sun committed
58
59
60
    if batch % self.log_steps == 0:
      elapsed_time = time.time() - self.start_time
      examples_per_second = (self._batch_size * self.log_steps) / elapsed_time
61
62
63
64
65
      self.record_batch = True
      # TODO(anjalisridhar): add timestamp as well.
      if batch != 0:
        tf.logging.info("BenchmarkMetric: {'num_batches':%d, 'time_taken': %f,"
                        "'images_per_second': %f}" %
Shining Sun's avatar
Shining Sun committed
66
                        (batch, elapsed_time, examples_per_second))
67

68

69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
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):
89
90
    if not hasattr(self.model.optimizer, 'learning_rate'):
      raise ValueError('Optimizer must have a "learning_rate" attribute.')
91
92
93
    self.epochs += 1

  def on_batch_begin(self, batch, logs=None):
94
95
96
97
    lr = self.schedule(self.epochs,
                       batch,
                       self.batches_per_epoch,
                       self.batch_size)
98
99
100
    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
101
      self.model.optimizer.learning_rate = lr  # lr should be a float here
102
      self.prev_lr = lr
103
104
105
      tf.logging.debug('Epoch %05d Batch %05d: LearningRateBatchScheduler '
                       'change learning rate to %s.', self.epochs, batch, lr)

106

Shining Sun's avatar
Shining Sun committed
107
def get_optimizer():
108
109
  """Returns optimizer to use."""
  # The learning_rate is overwritten at the beginning of each step by callback.
Shining Sun's avatar
Shining Sun committed
110
  return gradient_descent_v2.SGD(learning_rate=0.1, momentum=0.9)
111
112


113
def get_callbacks(learning_rate_schedule_fn, num_images):
114
  """Returns common callbacks."""
Shining Sun's avatar
Shining Sun committed
115
  time_callback = TimeHistory(FLAGS.batch_size)
116
117

  tensorboard_callback = tf.keras.callbacks.TensorBoard(
118
      log_dir=FLAGS.model_dir)
119

Shining Sun's avatar
Shining Sun committed
120
  lr_callback = LearningRateBatchScheduler(
121
122
123
      learning_rate_schedule_fn,
      batch_size=FLAGS.batch_size,
      num_images=num_images)
124
125
126

  return time_callback, tensorboard_callback, lr_callback

Shining Sun's avatar
Shining Sun committed
127

128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
def build_stats(history, eval_output):
  """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.

  Returns:
    Dictionary of normalized results.
  """
  stats = {}
  if eval_output:
    stats['accuracy_top_1'] = eval_output[1].item()
    stats['eval_loss'] = eval_output[0].item()
  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()

  return stats


Shining Sun's avatar
Shining Sun committed
157
158
def define_keras_flags():
  flags.DEFINE_boolean(name='enable_eager', default=False, help='Enable eager?')
159
  flags.DEFINE_boolean(name='skip_eval', default=False, help='Skip evaluation?')
Shining Sun's avatar
Shining Sun committed
160
  flags.DEFINE_integer(
161
162
163
164
      name='train_steps', default=None,
      help='The number of steps to run for training. If it is larger than '
      '# batches per epoch, then use # bathes per epoch. When this flag is '
      'set, only one epoch is going to run for training.')
165

Shining Sun's avatar
Shining Sun committed
166
167
168
169
170
171
172
173

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
174
  tuning the full input pipeline.
Shining Sun's avatar
Shining Sun committed
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209

  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].
    inputs = tf.truncated_normal(
        [batch_size] + [height, width, num_channels],
        dtype=dtype,
        mean=127,
        stddev=60,
        name='synthetic_inputs')

    labels = tf.random_uniform(
        [batch_size] + [1],
        minval=0,
        maxval=num_classes - 1,
        dtype=tf.int32,
        name='synthetic_labels')
    data = tf.data.Dataset.from_tensors((inputs, labels)).repeat()
    data = data.prefetch(buffer_size=tf.contrib.data.AUTOTUNE)
    return data

  return input_fn