cifar10_main.py 8.72 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

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

import os
22
import sys
23

Karmel Allison's avatar
Karmel Allison committed
24
import tensorflow as tf  # pylint: disable=g-bad-import-order
25

26
27
from official.resnet import resnet_model
from official.resnet import resnet_run_loop
28

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

38
39
40
41
_NUM_IMAGES = {
    'train': 50000,
    'validation': 10000,
}
42
43


44
45
46
###############################################################################
# Data processing
###############################################################################
47
def get_filenames(is_training, data_dir):
48
  """Returns a list of filenames."""
49
  data_dir = os.path.join(data_dir, 'cifar-10-batches-bin')
50

51
52
53
  assert os.path.exists(data_dir), (
      'Run cifar10_download_and_extract.py first to download and extract the '
      'CIFAR-10 data.')
54

55
  if is_training:
56
57
    return [
        os.path.join(data_dir, 'data_batch_%d.bin' % i)
58
        for i in range(1, _NUM_DATA_FILES + 1)
59
60
    ]
  else:
61
    return [os.path.join(data_dir, 'test_batch.bin')]
62
63


64
def parse_record(raw_record, is_training):
Kathy Wu's avatar
Kathy Wu committed
65
  """Parse CIFAR-10 image and label from a raw record."""
66
67
  # Convert bytes to a vector of uint8 that is record_bytes long.
  record_vector = tf.decode_raw(raw_record, tf.uint8)
68

69
70
  # The first byte represents the label, which we convert from uint8 to int32
  # and then to one-hot.
71
  label = tf.cast(record_vector[0], tf.int32)
72
  label = tf.one_hot(label, _NUM_CLASSES)
73
74
75

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

  # 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)

83
84
  image = preprocess_image(image, is_training)

85
  return image, label
86
87


88
89
90
91
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
92
93
    image = tf.image.resize_image_with_crop_or_pad(
        image, _HEIGHT + 8, _WIDTH + 8)
94

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

98
99
    # Randomly flip the image horizontally.
    image = tf.image.random_flip_left_right(image)
Kathy Wu's avatar
Kathy Wu committed
100
101
102

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


106
107
def input_fn(is_training, data_dir, batch_size, num_epochs=1,
             num_parallel_calls=1, multi_gpu=False):
108
  """Input_fn using the tf.data input pipeline for CIFAR-10 dataset.
109
110

  Args:
111
    is_training: A boolean denoting whether the input is for training.
Kathy Wu's avatar
Kathy Wu committed
112
    data_dir: The directory containing the input data.
113
    batch_size: The number of samples per batch.
114
    num_epochs: The number of epochs to repeat the dataset.
115
116
117
118
119
120
    num_parallel_calls: The number of records that are processed in parallel.
      This can be optimized per data set but for generally homogeneous data
      sets, should be approximately the number of available CPU cores.
    multi_gpu: Whether this is run multi-GPU. Note that this is only required
      currently to handle the batch leftovers, and can be removed
      when that is handled directly by Estimator.
121
122

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

128
129
  num_images = is_training and _NUM_IMAGES['train'] or _NUM_IMAGES['validation']

Karmel Allison's avatar
Karmel Allison committed
130
131
  return resnet_run_loop.process_record_dataset(
      dataset, is_training, batch_size, _NUM_IMAGES['train'],
132
133
      parse_record, num_epochs, num_parallel_calls,
      examples_per_epoch=num_images, multi_gpu=multi_gpu)
134
135


136
def get_synth_input_fn():
Karmel Allison's avatar
Karmel Allison committed
137
138
  return resnet_run_loop.get_synth_input_fn(
      _HEIGHT, _WIDTH, _NUM_CHANNELS, _NUM_CLASSES)
139
140


141
142
143
###############################################################################
# Running the model
###############################################################################
144
class Cifar10Model(resnet_model.Model):
Karmel Allison's avatar
Karmel Allison committed
145
  """Model class with appropriate defaults for CIFAR-10 data."""
146

147
  def __init__(self, resnet_size, data_format=None, num_classes=_NUM_CLASSES,
148
149
               version=resnet_model.DEFAULT_VERSION,
               dtype=resnet_model.DEFAULT_DTYPE):
Neal Wu's avatar
Neal Wu committed
150
151
152
153
154
155
156
    """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
157
        enables users to extend the same model to their own datasets.
158
159
      version: Integer representing which version of the ResNet network to use.
        See README for details. Valid values: [1, 2]
160
      dtype: The TensorFlow dtype to use for calculations.
Karmel Allison's avatar
Karmel Allison committed
161
162
163

    Raises:
      ValueError: if invalid resnet_size is chosen
Neal Wu's avatar
Neal Wu committed
164
    """
165
166
167
168
169
170
171
    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,
172
        bottleneck=False,
173
        num_classes=num_classes,
174
175
176
177
178
179
180
181
182
183
        num_filters=16,
        kernel_size=3,
        conv_stride=1,
        first_pool_size=None,
        first_pool_stride=None,
        second_pool_size=8,
        second_pool_stride=1,
        block_sizes=[num_blocks] * 3,
        block_strides=[1, 2, 2],
        final_size=64,
184
        version=version,
185
186
187
        data_format=data_format,
        dtype=dtype
    )
188
189


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

194
  learning_rate_fn = resnet_run_loop.learning_rate_with_decay(
195
196
197
198
199
200
201
202
203
204
205
206
207
      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
208
  def loss_filter_fn(_):
209
210
    return True

211
212
213
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'],
      version=params['version'],
      loss_scale=params['loss_scale'],
      loss_filter_fn=loss_filter_fn,
224
      multi_gpu=params['multi_gpu'],
225
226
      dtype=params['dtype']
  )
227
228


229
def main(argv):
230
  parser = resnet_run_loop.ResnetArgParser()
231
232
233
234
235
  # Set defaults that are reasonable for this model.
  parser.set_defaults(data_dir='/tmp/cifar10_data',
                      model_dir='/tmp/cifar10_model',
                      resnet_size=32,
                      train_epochs=250,
236
                      epochs_between_evals=10,
237
238
                      batch_size=128)

239
240
241
  flags = parser.parse_args(args=argv[1:])

  input_function = flags.use_synthetic_data and get_synth_input_fn() or input_fn
242
243
244
245

  resnet_run_loop.resnet_main(
      flags, cifar10_model_fn, input_function,
      shape=[_HEIGHT, _WIDTH, _NUM_CHANNELS])
246
247
248
249


if __name__ == '__main__':
  tf.logging.set_verbosity(tf.logging.INFO)
250
  main(argv=sys.argv)