cifar10_main.py 8.92 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
# Copyright 2017 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
"""Runs a ResNet model on the CIFAR-10 dataset."""
16
17
18
19
20
21
22

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

import os

23
24
from absl import app as absl_app
from absl import flags
Karmel Allison's avatar
Karmel Allison committed
25
import tensorflow as tf  # pylint: disable=g-bad-import-order
26

27
from official.utils.flags import core as flags_core
28
from official.utils.logs import logger
29
30
from official.resnet import resnet_model
from official.resnet import resnet_run_loop
31

32
33
_HEIGHT = 32
_WIDTH = 32
34
35
_NUM_CHANNELS = 3
_DEFAULT_IMAGE_BYTES = _HEIGHT * _WIDTH * _NUM_CHANNELS
36
37
# The record is the image plus a one-byte label
_RECORD_BYTES = _DEFAULT_IMAGE_BYTES + 1
38
39
40
_NUM_CLASSES = 10
_NUM_DATA_FILES = 5

41
42
43
44
_NUM_IMAGES = {
    'train': 50000,
    'validation': 10000,
}
45

46
47
DATASET_NAME = 'CIFAR-10'

48

49
50
51
###############################################################################
# Data processing
###############################################################################
52
def get_filenames(is_training, data_dir):
53
  """Returns a list of filenames."""
54
  data_dir = os.path.join(data_dir, 'cifar-10-batches-bin')
55

56
57
58
  assert os.path.exists(data_dir), (
      'Run cifar10_download_and_extract.py first to download and extract the '
      'CIFAR-10 data.')
59

60
  if is_training:
61
62
    return [
        os.path.join(data_dir, 'data_batch_%d.bin' % i)
63
        for i in range(1, _NUM_DATA_FILES + 1)
64
65
    ]
  else:
66
    return [os.path.join(data_dir, 'test_batch.bin')]
67
68


69
def parse_record(raw_record, is_training, dtype):
Kathy Wu's avatar
Kathy Wu committed
70
  """Parse CIFAR-10 image and label from a raw record."""
71
72
  # Convert bytes to a vector of uint8 that is record_bytes long.
  record_vector = tf.decode_raw(raw_record, tf.uint8)
73

74
75
  # The first byte represents the label, which we convert from uint8 to int32
  # and then to one-hot.
76
  label = tf.cast(record_vector[0], tf.int32)
77
78
79

  # The remaining bytes after the label represent the image, which we reshape
  # from [depth * height * width] to [depth, height, width].
80
  depth_major = tf.reshape(record_vector[1:_RECORD_BYTES],
81
                           [_NUM_CHANNELS, _HEIGHT, _WIDTH])
82
83
84
85
86

  # Convert from [depth, height, width] to [height, width, depth], and cast as
  # float32.
  image = tf.cast(tf.transpose(depth_major, [1, 2, 0]), tf.float32)

87
  image = preprocess_image(image, is_training)
88
  image = tf.cast(image, dtype)
89

90
  return image, label
91
92


93
94
95
96
def preprocess_image(image, is_training):
  """Preprocess a single image of layout [height, width, depth]."""
  if is_training:
    # Resize the image to add four extra pixels on each side.
Neal Wu's avatar
Neal Wu committed
97
98
    image = tf.image.resize_image_with_crop_or_pad(
        image, _HEIGHT + 8, _WIDTH + 8)
99

100
    # Randomly crop a [_HEIGHT, _WIDTH] section of the image.
101
    image = tf.random_crop(image, [_HEIGHT, _WIDTH, _NUM_CHANNELS])
Kathy Wu's avatar
Kathy Wu committed
102

103
104
    # Randomly flip the image horizontally.
    image = tf.image.random_flip_left_right(image)
Kathy Wu's avatar
Kathy Wu committed
105
106
107

  # Subtract off the mean and divide by the variance of the pixels.
  image = tf.image.per_image_standardization(image)
108
  return image
Kathy Wu's avatar
Kathy Wu committed
109
110


111
112
113
def input_fn(is_training, data_dir, batch_size, num_epochs=1, num_gpus=None,
             dtype=tf.float32):
  """Input function which provides batches for train or eval.
114
115

  Args:
116
    is_training: A boolean denoting whether the input is for training.
Kathy Wu's avatar
Kathy Wu committed
117
    data_dir: The directory containing the input data.
118
    batch_size: The number of samples per batch.
119
    num_epochs: The number of epochs to repeat the dataset.
Taylor Robie's avatar
Taylor Robie committed
120
    num_gpus: The number of gpus used for training.
121
    dtype: Data type to use for images/features
122
123

  Returns:
124
    A dataset that can be used for iteration.
125
  """
126
127
  filenames = get_filenames(is_training, data_dir)
  dataset = tf.data.FixedLengthRecordDataset(filenames, _RECORD_BYTES)
128

Karmel Allison's avatar
Karmel Allison committed
129
  return resnet_run_loop.process_record_dataset(
Taylor Robie's avatar
Taylor Robie committed
130
131
132
133
134
135
136
      dataset=dataset,
      is_training=is_training,
      batch_size=batch_size,
      shuffle_buffer=_NUM_IMAGES['train'],
      parse_record_fn=parse_record,
      num_epochs=num_epochs,
      num_gpus=num_gpus,
137
138
      examples_per_epoch=_NUM_IMAGES['train'] if is_training else None,
      dtype=dtype
139
  )
140
141


Toby Boyd's avatar
Toby Boyd committed
142
def get_synth_input_fn(dtype):
Karmel Allison's avatar
Karmel Allison committed
143
  return resnet_run_loop.get_synth_input_fn(
Toby Boyd's avatar
Toby Boyd committed
144
      _HEIGHT, _WIDTH, _NUM_CHANNELS, _NUM_CLASSES, dtype=dtype)
145
146


147
148
149
###############################################################################
# Running the model
###############################################################################
150
class Cifar10Model(resnet_model.Model):
Karmel Allison's avatar
Karmel Allison committed
151
  """Model class with appropriate defaults for CIFAR-10 data."""
152

153
  def __init__(self, resnet_size, data_format=None, num_classes=_NUM_CLASSES,
154
               resnet_version=resnet_model.DEFAULT_VERSION,
155
               dtype=resnet_model.DEFAULT_DTYPE):
Neal Wu's avatar
Neal Wu committed
156
157
158
159
160
161
162
    """These are the parameters that work for CIFAR-10 data.

    Args:
      resnet_size: The number of convolutional layers needed in the model.
      data_format: Either 'channels_first' or 'channels_last', specifying which
        data format to use when setting up the model.
      num_classes: The number of output classes needed from the model. This
163
        enables users to extend the same model to their own datasets.
164
165
      resnet_version: Integer representing which version of the ResNet network
      to use. See README for details. Valid values: [1, 2]
166
      dtype: The TensorFlow dtype to use for calculations.
Karmel Allison's avatar
Karmel Allison committed
167
168
169

    Raises:
      ValueError: if invalid resnet_size is chosen
Neal Wu's avatar
Neal Wu committed
170
    """
171
172
173
174
175
176
177
    if resnet_size % 6 != 2:
      raise ValueError('resnet_size must be 6n + 2:', resnet_size)

    num_blocks = (resnet_size - 2) // 6

    super(Cifar10Model, self).__init__(
        resnet_size=resnet_size,
178
        bottleneck=False,
179
        num_classes=num_classes,
180
181
182
183
184
185
186
        num_filters=16,
        kernel_size=3,
        conv_stride=1,
        first_pool_size=None,
        first_pool_stride=None,
        block_sizes=[num_blocks] * 3,
        block_strides=[1, 2, 2],
187
        resnet_version=resnet_version,
188
189
190
        data_format=data_format,
        dtype=dtype
    )
191
192


193
194
195
196
def cifar10_model_fn(features, labels, mode, params):
  """Model function for CIFAR-10."""
  features = tf.reshape(features, [-1, _HEIGHT, _WIDTH, _NUM_CHANNELS])

197
  learning_rate_fn = resnet_run_loop.learning_rate_with_decay(
198
199
200
201
202
203
204
205
206
207
208
209
210
      batch_size=params['batch_size'], batch_denom=128,
      num_images=_NUM_IMAGES['train'], boundary_epochs=[100, 150, 200],
      decay_rates=[1, 0.1, 0.01, 0.001])

  # We use a weight decay of 0.0002, which performs better
  # than the 0.0001 that was originally suggested.
  weight_decay = 2e-4

  # Empirical testing showed that including batch_normalization variables
  # in the calculation of regularized loss helped validation accuracy
  # for the CIFAR-10 dataset, perhaps because the regularization prevents
  # overfitting on the small data set. We therefore include all vars when
  # regularizing and computing loss during training.
Karmel Allison's avatar
Karmel Allison committed
211
  def loss_filter_fn(_):
212
213
    return True

214
215
216
217
218
219
220
221
222
223
  return resnet_run_loop.resnet_model_fn(
      features=features,
      labels=labels,
      mode=mode,
      model_class=Cifar10Model,
      resnet_size=params['resnet_size'],
      weight_decay=weight_decay,
      learning_rate_fn=learning_rate_fn,
      momentum=0.9,
      data_format=params['data_format'],
224
      resnet_version=params['resnet_version'],
225
226
      loss_scale=params['loss_scale'],
      loss_filter_fn=loss_filter_fn,
Zac Wellmer's avatar
Zac Wellmer committed
227
228
      dtype=params['dtype'],
      fine_tune=params['fine_tune']
229
  )
230
231


232
233
234
235
236
237
238
239
240
def define_cifar_flags():
  resnet_run_loop.define_resnet_flags()
  flags.adopt_module_key_flags(resnet_run_loop)
  flags_core.set_defaults(data_dir='/tmp/cifar10_data',
                          model_dir='/tmp/cifar10_model',
                          resnet_size='32',
                          train_epochs=250,
                          epochs_between_evals=10,
                          batch_size=128)
241

242

243
244
245
246
247
248
def run_cifar(flags_obj):
  """Run ResNet CIFAR-10 training and eval loop.

  Args:
    flags_obj: An object containing parsed flag values.
  """
Toby Boyd's avatar
Toby Boyd committed
249
250
251
  input_function = (flags_obj.use_synthetic_data and
                    get_synth_input_fn(flags_core.get_tf_dtype(flags_obj)) or
                    input_fn)
252
  resnet_run_loop.resnet_main(
253
      flags_obj, cifar10_model_fn, input_function, DATASET_NAME,
254
      shape=[_HEIGHT, _WIDTH, _NUM_CHANNELS])
255
256


257
def main(_):
258
259
  with logger.benchmark_context(flags.FLAGS):
    run_cifar(flags.FLAGS)
260
261


262
263
if __name__ == '__main__':
  tf.logging.set_verbosity(tf.logging.INFO)
264
265
  define_cifar_flags()
  absl_app.run(main)