keras_cifar_main.py 6.75 KB
Newer Older
Shining Sun's avatar
Shining Sun committed
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.
# ==============================================================================
Shining Sun's avatar
Shining Sun committed
15
"""Runs a ResNet model on the Cifar-10 dataset."""
16
17
18
19
20
21
22
23
24
25

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

from absl import app as absl_app
from absl import flags
import tensorflow as tf  # pylint: disable=g-bad-import-order

from official.resnet import cifar10_main as cifar_main
26
from official.resnet.keras import keras_common
Shining Sun's avatar
Shining Sun committed
27
from official.resnet.keras import resnet_cifar_model
28
29
30
31
32
from official.utils.flags import core as flags_core
from official.utils.logs import logger
from official.utils.misc import distribution_utils


33
34
LR_SCHEDULE = [  # (multiplier, epoch to start) tuples
    (0.1, 91), (0.01, 136), (0.001, 182)
35
36
]

37

38
39
40
41
def learning_rate_schedule(current_epoch,
                           current_batch,
                           batches_per_epoch,
                           batch_size):
Shining Sun's avatar
Shining Sun committed
42
  """Handles linear scaling rule and LR decay.
43

44
45
  Scale learning rate at epoch boundaries provided in LR_SCHEDULE by the
  provided scaling factor.
46
47
48
49

  Args:
    current_epoch: integer, current epoch indexed from 0.
    current_batch: integer, current batch in the current epoch, indexed from 0.
50
51
    batches_per_epoch: integer, number of steps in an epoch.
    batch_size: integer, total batch sized.
52
53
54
55

  Returns:
    Adjusted learning rate.
  """
Shining Sun's avatar
Shining Sun committed
56
  initial_learning_rate = keras_common.BASE_LEARNING_RATE * batch_size / 128
57
  learning_rate = initial_learning_rate
58
  for mult, start_epoch in LR_SCHEDULE:
59
60
    if current_epoch >= start_epoch:
      learning_rate = initial_learning_rate * mult
61
62
63
64
65
66
67
68
69
70
71
    else:
      break
  return learning_rate


def parse_record_keras(raw_record, is_training, dtype):
  """Parses a record containing a training example of an image.

  The input record is parsed into a label and image, and the image is passed
  through preprocessing steps (cropping, flipping, and so on).

Shining Sun's avatar
Shining Sun committed
72
  This method converts the label to one hot to fit the loss function.
73

74
75
76
77
78
79
80
81
82
83
  Args:
    raw_record: scalar Tensor tf.string containing a serialized
      Example protocol buffer.
    is_training: A boolean denoting whether the input is for training.
    dtype: Data type to use for input images.

  Returns:
    Tuple with processed image tensor and one-hot-encoded label tensor.
  """
  image, label = cifar_main.parse_record(raw_record, is_training, dtype)
84
  label = tf.compat.v1.sparse_to_dense(label, (cifar_main.NUM_CLASSES,), 1)
85
86
87
  return image, label


Shining Sun's avatar
Shining Sun committed
88
89
def run(flags_obj):
  """Run ResNet Cifar-10 training and eval loop using native Keras APIs.
90
91
92
93
94
95

  Args:
    flags_obj: An object containing parsed flag values.

  Raises:
    ValueError: If fp16 is passed as it is not currently supported.
96
97
98

  Returns:
    Dictionary of training and eval stats.
99
  """
100
  if flags_obj.enable_eager:
101
    tf.compat.v1.enable_eager_execution()
102

103
104
105
106
107
  dtype = flags_core.get_tf_dtype(flags_obj)
  if dtype == 'fp16':
    raise ValueError('dtype fp16 is not supported in Keras. Use the default '
                     'value(fp32).')

108
109
110
111
112
  data_format = flags_obj.data_format
  if data_format is None:
    data_format = ('channels_first'
                   if tf.test.is_built_with_cuda() else 'channels_last')
  tf.keras.backend.set_image_data_format(data_format)
113

114
  if flags_obj.use_synthetic_data:
115
    distribution_utils.set_up_synthetic_data()
Shining Sun's avatar
Shining Sun committed
116
    input_fn = keras_common.get_synth_input_fn(
117
118
119
120
        height=cifar_main.HEIGHT,
        width=cifar_main.WIDTH,
        num_channels=cifar_main.NUM_CHANNELS,
        num_classes=cifar_main.NUM_CLASSES,
Shining Sun's avatar
Shining Sun committed
121
        dtype=flags_core.get_tf_dtype(flags_obj))
122
  else:
123
    distribution_utils.undo_set_up_synthetic_data()
Shining Sun's avatar
Shining Sun committed
124
125
126
127
128
    input_fn = cifar_main.input_fn

  train_input_dataset = input_fn(
      is_training=True,
      data_dir=flags_obj.data_dir,
129
      batch_size=flags_obj.batch_size,
Shining Sun's avatar
Shining Sun committed
130
131
132
133
134
135
      num_epochs=flags_obj.train_epochs,
      parse_record_fn=parse_record_keras)

  eval_input_dataset = input_fn(
      is_training=False,
      data_dir=flags_obj.data_dir,
136
      batch_size=flags_obj.batch_size,
Shining Sun's avatar
Shining Sun committed
137
138
      num_epochs=flags_obj.train_epochs,
      parse_record_fn=parse_record_keras)
139

140
  strategy = distribution_utils.get_distribution_strategy(
Shining Sun's avatar
Shining Sun committed
141
142
      num_gpus=flags_obj.num_gpus,
      turn_off_distribution_strategy=flags_obj.turn_off_distribution_strategy)
143

Shining Sun's avatar
Shining Sun committed
144
145
146
  strategy_scope = keras_common.get_strategy_scope(strategy)

  with strategy_scope:
Shining Sun's avatar
Shining Sun committed
147
148
    optimizer = keras_common.get_optimizer()
    model = resnet_cifar_model.resnet56(classes=cifar_main.NUM_CLASSES)
Shining Sun's avatar
Shining Sun committed
149

Shining Sun's avatar
Shining Sun committed
150
151
152
    model.compile(loss='categorical_crossentropy',
                  optimizer=optimizer,
                  metrics=['categorical_accuracy'])
Shining Sun's avatar
Shining Sun committed
153

154
155
  time_callback, tensorboard_callback, lr_callback = keras_common.get_callbacks(
      learning_rate_schedule, cifar_main.NUM_IMAGES['train'])
156

Shining Sun's avatar
Shining Sun committed
157
158
159
160
161
162
163
  train_steps = cifar_main.NUM_IMAGES['train'] // flags_obj.batch_size
  train_epochs = flags_obj.train_epochs

  if flags_obj.train_steps:
    train_steps = min(flags_obj.train_steps, train_steps)
    train_epochs = 1

164
  num_eval_steps = (cifar_main.NUM_IMAGES['validation'] //
165
166
                    flags_obj.batch_size)

Shining Sun's avatar
Shining Sun committed
167
168
  validation_data = eval_input_dataset
  if flags_obj.skip_eval:
169
    tf.keras.backend.set_learning_phase(1)
Shining Sun's avatar
Shining Sun committed
170
171
172
    num_eval_steps = None
    validation_data = None

173
  history = model.fit(train_input_dataset,
174
175
176
177
178
179
180
181
182
                      epochs=train_epochs,
                      steps_per_epoch=train_steps,
                      callbacks=[
                          time_callback,
                          lr_callback,
                          tensorboard_callback
                      ],
                      validation_steps=num_eval_steps,
                      validation_data=validation_data,
183
                      verbose=2)
184
  eval_output = None
185
  if not flags_obj.skip_eval:
Shining Sun's avatar
Shining Sun committed
186
187
188
    eval_output = model.evaluate(eval_input_dataset,
                                 steps=num_eval_steps,
                                 verbose=1)
189
  stats = keras_common.build_stats(history, eval_output, time_callback)
190
  return stats
191

192
193

def main(_):
194
  with logger.benchmark_context(flags.FLAGS):
195
    return run(flags.FLAGS)
196
197
198


if __name__ == '__main__':
199
  tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.INFO)
200
  cifar_main.define_cifar_flags()
Shining Sun's avatar
Shining Sun committed
201
  keras_common.define_keras_flags()
202
  absl_app.run(main)