keras_imagenet_main.py 7.45 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
15
16
17
18
19
20
21
22
23
24
25
#
# 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.
# ==============================================================================
"""Runs a ResNet model on the ImageNet dataset."""

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 imagenet_main
26
from official.resnet import imagenet_preprocessing
27
from official.resnet import resnet_run_loop
28
from official.resnet.keras import keras_common
Shining Sun's avatar
Shining Sun committed
29
from official.resnet.keras import resnet50
30
31
32
33
from official.utils.flags import core as flags_core
from official.utils.logs import logger
from official.utils.misc import distribution_utils

34
35
# import os
# os.environ['TF2_BEHAVIOR'] = 'enabled'
36
37
38
39

LR_SCHEDULE = [    # (multiplier, epoch to start) tuples
    (1.0, 5), (0.1, 30), (0.01, 60), (0.001, 80)
]
Shining Sun's avatar
Shining Sun committed
40

41

42
def learning_rate_schedule(current_epoch, current_batch, batches_per_epoch, batch_size):
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
  """Handles linear scaling rule, gradual warmup, and LR decay.

  The learning rate starts at 0, then it increases linearly per step.
  After 5 epochs we reach the base learning rate (scaled to account
    for batch size).
  After 30, 60 and 80 epochs the learning rate is divided by 10.
  After 90 epochs training stops and the LR is set to 0. This ensures
    that we train for exactly 90 epochs for reproducibility.

  Args:
    current_epoch: integer, current epoch indexed from 0.
    current_batch: integer, current batch in the current epoch, indexed from 0.

  Returns:
    Adjusted learning rate.
  """
Shining Sun's avatar
Shining Sun committed
59
  initial_learning_rate = keras_common.BASE_LEARNING_RATE * batch_size / 256
60
61
62
63
  epoch = current_epoch + float(current_batch) / batches_per_epoch
  warmup_lr_multiplier, warmup_end_epoch = LR_SCHEDULE[0]
  if epoch < warmup_end_epoch:
    # Learning rate increases linearly per step.
64
    return initial_learning_rate * warmup_lr_multiplier * epoch / warmup_end_epoch
65
66
  for mult, start_epoch in LR_SCHEDULE:
    if epoch >= start_epoch:
67
      learning_rate = initial_learning_rate * mult
68
69
70
71
72
73
    else:
      break
  return learning_rate


def parse_record_keras(raw_record, is_training, dtype):
Shining Sun's avatar
Shining Sun committed
74
  """Adjust the shape of label."""
75
76
77
78
79
80
81
82
83
84
85
86
87
  image_buffer, label, bbox = imagenet_main._parse_example_proto(raw_record)

  image = imagenet_preprocessing.preprocess_image(
      image_buffer=image_buffer,
      bbox=bbox,
      output_height=imagenet_main.DEFAULT_IMAGE_SIZE,
      output_width=imagenet_main.DEFAULT_IMAGE_SIZE,
      num_channels=imagenet_main.NUM_CHANNELS,
      is_training=is_training)
  image = tf.cast(image, dtype)
  label = tf.sparse_to_dense(label, (imagenet_main.NUM_CLASSES,), 1)

  """
Shining Sun's avatar
Shining Sun committed
88
89
90
91
  image, label = imagenet_main.parse_record(raw_record, is_training, dtype)
  # Subtract one so that labels are in [0, 1000), and cast to float32 for
  # Keras model.
  label = tf.cast(tf.cast(tf.reshape(label, shape=[1]), dtype=tf.int32) - 1,
Shining Sun's avatar
Shining Sun committed
92
      dtype=tf.float32)
93
  """
94
95
96
97
98
99
100
101
102
103
104
105
  return image, label


def run_imagenet_with_keras(flags_obj):
  """Run ResNet ImageNet training and eval loop using native Keras APIs.

  Args:
    flags_obj: An object containing parsed flag values.

  Raises:
    ValueError: If fp16 is passed as it is not currently supported.
  """
106
107
  if flags_obj.enable_eager:
    tf.enable_eager_execution()
Shining Sun's avatar
Shining Sun committed
108

109
110
111
112
113
114
115
116
117
118
119
  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).')

  per_device_batch_size = distribution_utils.per_device_batch_size(
      flags_obj.batch_size, flags_core.get_num_gpus(flags_obj))

  # pylint: disable=protected-access
  if flags_obj.use_synthetic_data:
    synth_input_fn = resnet_run_loop.get_synth_input_fn(
120
121
        imagenet_main.DEFAULT_IMAGE_SIZE, imagenet_main.DEFAULT_IMAGE_SIZE,
        imagenet_main.NUM_CHANNELS, imagenet_main.NUM_CLASSES,
122
123
124
        dtype=flags_core.get_tf_dtype(flags_obj))
    train_input_dataset = synth_input_fn(
        batch_size=per_device_batch_size,
125
126
127
128
        height=imagenet_main.DEFAULT_IMAGE_SIZE,
        width=imagenet_main.DEFAULT_IMAGE_SIZE,
        num_channels=imagenet_main.NUM_CHANNELS,
        num_classes=imagenet_main.NUM_CLASSES,
129
130
131
        dtype=dtype)
    eval_input_dataset = synth_input_fn(
        batch_size=per_device_batch_size,
132
133
134
135
        height=imagenet_main.DEFAULT_IMAGE_SIZE,
        width=imagenet_main.DEFAULT_IMAGE_SIZE,
        num_channels=imagenet_main.NUM_CHANNELS,
        num_classes=imagenet_main.NUM_CLASSES,
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
        dtype=dtype)
  # pylint: enable=protected-access

  else:
    train_input_dataset = imagenet_main.input_fn(
          True,
          flags_obj.data_dir,
          batch_size=per_device_batch_size,
          num_epochs=flags_obj.train_epochs,
          parse_record_fn=parse_record_keras)

    eval_input_dataset = imagenet_main.input_fn(
          False,
          flags_obj.data_dir,
          batch_size=per_device_batch_size,
          num_epochs=flags_obj.train_epochs,
          parse_record_fn=parse_record_keras)


Shining Sun's avatar
Shining Sun committed
155
  optimizer = keras_common.get_optimizer()
156
157
  strategy = distribution_utils.get_distribution_strategy(
    flags_obj.num_gpus, flags_obj.use_one_device_strategy)
158

159
  model = resnet50.ResNet50(num_classes=imagenet_main.NUM_CLASSES)
Shining Sun's avatar
Shining Sun committed
160

Shining Sun's avatar
Shining Sun committed
161
162
163
  model.compile(loss='categorical_crossentropy',
                optimizer=optimizer,
                metrics=['categorical_accuracy'],
164
                distribute=strategy)
Shining Sun's avatar
Shining Sun committed
165

166
167
  time_callback, tensorboard_callback, lr_callback = keras_common.get_callbacks(
      learning_rate_schedule, imagenet_main.NUM_IMAGES['train'])
168

169
170
  steps_per_epoch = imagenet_main.NUM_IMAGES['train'] // flags_obj.batch_size
  num_eval_steps = (imagenet_main.NUM_IMAGES['validation'] //
171
                  flags_obj.batch_size)
Shining Sun's avatar
Shining Sun committed
172

Shining Sun's avatar
Shining Sun committed
173
174
175
176
177
178
179
  train_steps = imagenet_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

Shining Sun's avatar
bug fix  
Shining Sun committed
180
  history = model.fit(train_input_dataset,
Shining Sun's avatar
Shining Sun committed
181
182
                      epochs=train_epochs,
                      steps_per_epoch=train_steps,
Shining Sun's avatar
bug fix  
Shining Sun committed
183
184
185
186
187
188
189
190
                      callbacks=[
                        time_callback,
                        lr_callback,
                        tensorboard_callback
                      ],
                      validation_steps=num_eval_steps,
                      validation_data=eval_input_dataset,
                      verbose=1)
Shining Sun's avatar
Shining Sun committed
191

192
193
194
195
  if not flags_obj.skip_eval:
    eval_output = model.evaluate(eval_input_dataset,
                                 steps=num_eval_steps,
                                 verbose=1)
Shining Sun's avatar
bug fix  
Shining Sun committed
196
197

  stats = keras_common.analyze_fit_and_eval_result(history, eval_output)
198
199

  return stats
200

Shining Sun's avatar
Shining Sun committed
201

202
203
204
205
206
207
def main(_):
  with logger.benchmark_context(flags.FLAGS):
    run_imagenet_with_keras(flags.FLAGS)


if __name__ == '__main__':
208
  tf.logging.set_verbosity(tf.logging.INFO)
209
  imagenet_main.define_imagenet_flags()
Shining Sun's avatar
Shining Sun committed
210
  keras_common.define_keras_flags()
211
  absl_app.run(main)