cifar10_main.py 7.5 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
24
25

import tensorflow as tf

Karmel Allison's avatar
Karmel Allison committed
26
import resnet
27

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

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


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

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

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


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

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

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

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

82
83
  image = preprocess_image(image, is_training)

84
  return image, label
85
86


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

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

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

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


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

  Args:
110
    is_training: A boolean denoting whether the input is for training.
Kathy Wu's avatar
Kathy Wu committed
111
    data_dir: The directory containing the input data.
112
    batch_size: The number of samples per batch.
113
    num_epochs: The number of epochs to repeat the dataset.
114
115
116
    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.
117
118

  Returns:
119
    A dataset that can be used for iteration.
120
  """
121
122
  filenames = get_filenames(is_training, data_dir)
  dataset = tf.data.FixedLengthRecordDataset(filenames, _RECORD_BYTES)
123

124
125
  return resnet.process_record_dataset(dataset, is_training, batch_size,
      _NUM_IMAGES['train'], parse_record, num_epochs, num_parallel_calls)
126
127


128
129
130
###############################################################################
# Running the model
###############################################################################
Karmel Allison's avatar
Karmel Allison committed
131
class Cifar10Model(resnet.Model):
132
133

  def __init__(self, resnet_size, data_format=None, num_classes=_NUM_CLASSES):
Neal Wu's avatar
Neal Wu committed
134
135
136
137
138
139
140
141
142
    """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
        enables users to extend the same model to different datasets.
    """
143
144
145
146
147
148
149
    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,
150
        num_classes=num_classes,
151
152
153
154
155
156
157
        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,
Karmel Allison's avatar
Karmel Allison committed
158
        block_fn=resnet.building_block,
159
160
161
162
        block_sizes=[num_blocks] * 3,
        block_strides=[1, 2, 2],
        final_size=64,
        data_format=data_format)
163
164


165
166
167
168
def cifar10_model_fn(features, labels, mode, params):
  """Model function for CIFAR-10."""
  features = tf.reshape(features, [-1, _HEIGHT, _WIDTH, _NUM_CHANNELS])

Karmel Allison's avatar
Karmel Allison committed
169
  learning_rate_fn = resnet.learning_rate_with_decay(
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
      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.
  def loss_filter_fn(name):
    return True

Karmel Allison's avatar
Karmel Allison committed
186
187
188
189
190
191
192
  return resnet.resnet_model_fn(features, labels, mode, Cifar10Model,
                                resnet_size=params['resnet_size'],
                                weight_decay=weight_decay,
                                learning_rate_fn=learning_rate_fn,
                                momentum=0.9,
                                data_format=params['data_format'],
                                loss_filter_fn=loss_filter_fn)
193
194
195


def main(unused_argv):
Karmel Allison's avatar
Karmel Allison committed
196
  resnet.resnet_main(FLAGS, cifar10_model_fn, input_fn)
197
198
199
200


if __name__ == '__main__':
  tf.logging.set_verbosity(tf.logging.INFO)
201

Karmel Allison's avatar
Karmel Allison committed
202
  parser = resnet.ResnetArgParser()
203
204
205
206
207
208
209
210
  # 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,
                      epochs_per_eval=10,
                      batch_size=128)

211
212
  FLAGS, unparsed = parser.parse_known_args()
  tf.app.run(argv=[sys.argv[0]] + unparsed)