cifar10_main.py 10.3 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 argparse
import os
23
import sys
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40

import tensorflow as tf

import resnet_model

parser = argparse.ArgumentParser()

# Basic model parameters.
parser.add_argument('--data_dir', type=str, default='/tmp/cifar10_data',
                    help='The path to the CIFAR-10 data directory.')

parser.add_argument('--model_dir', type=str, default='/tmp/cifar10_model',
                    help='The directory where the model will be stored.')

parser.add_argument('--resnet_size', type=int, default=32,
                    help='The size of the ResNet model to use.')

41
42
parser.add_argument('--train_epochs', type=int, default=250,
                    help='The number of epochs to train.')
43

44
parser.add_argument('--epochs_per_eval', type=int, default=10,
45
                    help='The number of epochs to run in between evaluations.')
46
47
48
49

parser.add_argument('--batch_size', type=int, default=128,
                    help='The number of images per batch.')

50
51
52
53
54
55
56
57
parser.add_argument(
    '--data_format', type=str, default=None,
    choices=['channels_first', 'channels_last'],
    help='A flag to override the data format used in the model. channels_first '
         'provides a performance boost on GPU but is not always compatible '
         'with CPU. If left unspecified, the data format will be chosen '
         'automatically based on whether TensorFlow was built for CPU or GPU.')

58
59
60
61
62
63
_HEIGHT = 32
_WIDTH = 32
_DEPTH = 3
_NUM_CLASSES = 10
_NUM_DATA_FILES = 5

64
65
66
# We use a weight decay of 0.0002, which performs better than the 0.0001 that
# was originally suggested.
_WEIGHT_DECAY = 2e-4
67
_MOMENTUM = 0.9
68

69
70
71
72
_NUM_IMAGES = {
    'train': 50000,
    'validation': 10000,
}
73
74
75
76


def record_dataset(filenames):
  """Returns an input pipeline Dataset from `filenames`."""
77
  record_bytes = _HEIGHT * _WIDTH * _DEPTH + 1
78
  return tf.data.FixedLengthRecordDataset(filenames, record_bytes)
79
80


81
def get_filenames(is_training, data_dir):
82
  """Returns a list of filenames."""
83
  data_dir = os.path.join(data_dir, 'cifar-10-batches-bin')
84

85
86
87
  assert os.path.exists(data_dir), (
      'Run cifar10_download_and_extract.py first to download and extract the '
      'CIFAR-10 data.')
88

89
  if is_training:
90
91
    return [
        os.path.join(data_dir, 'data_batch_%d.bin' % i)
92
        for i in range(1, _NUM_DATA_FILES + 1)
93
94
    ]
  else:
95
    return [os.path.join(data_dir, 'test_batch.bin')]
96
97


Kathy Wu's avatar
Kathy Wu committed
98
99
def parse_record(raw_record):
  """Parse CIFAR-10 image and label from a raw record."""
100
101
102
  # Every record consists of a label followed by the image, with a fixed number
  # of bytes for each.
  label_bytes = 1
103
  image_bytes = _HEIGHT * _WIDTH * _DEPTH
104
105
  record_bytes = label_bytes + image_bytes

106
107
  # Convert bytes to a vector of uint8 that is record_bytes long.
  record_vector = tf.decode_raw(raw_record, tf.uint8)
108

109
110
  # The first byte represents the label, which we convert from uint8 to int32
  # and then to one-hot.
111
  label = tf.cast(record_vector[0], tf.int32)
112
  label = tf.one_hot(label, _NUM_CLASSES)
113
114
115

  # The remaining bytes after the label represent the image, which we reshape
  # from [depth * height * width] to [depth, height, width].
116
117
  depth_major = tf.reshape(
      record_vector[label_bytes:record_bytes], [_DEPTH, _HEIGHT, _WIDTH])
118
119
120
121
122

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

123
  return image, label
124
125


126
127
128
129
130
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.
    image = tf.image.resize_image_with_crop_or_pad(image, _HEIGHT + 8, _WIDTH + 8)
131

132
133
    # Randomly crop a [_HEIGHT, _WIDTH] section of the image.
    image = tf.random_crop(image, [_HEIGHT, _WIDTH, _DEPTH])
Kathy Wu's avatar
Kathy Wu committed
134

135
136
    # Randomly flip the image horizontally.
    image = tf.image.random_flip_left_right(image)
Kathy Wu's avatar
Kathy Wu committed
137
138
139

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


143
def input_fn(is_training, data_dir, batch_size, num_epochs=1):
144
  """Input_fn using the tf.data input pipeline for CIFAR-10 dataset.
145
146

  Args:
147
    is_training: A boolean denoting whether the input is for training.
Kathy Wu's avatar
Kathy Wu committed
148
    data_dir: The directory containing the input data.
149
    batch_size: The number of samples per batch.
150
    num_epochs: The number of epochs to repeat the dataset.
151
152
153

  Returns:
    A tuple of images and labels.
154
  """
155
  dataset = record_dataset(get_filenames(is_training, data_dir))
156

157
  if is_training:
158
    # When choosing shuffle buffer sizes, larger sizes result in better
159
160
161
    # randomness, while smaller sizes have better performance. Because CIFAR-10
    # is a relatively small dataset, we choose to shuffle the full epoch.
    dataset = dataset.shuffle(buffer_size=_NUM_IMAGES['train'])
162

163
  dataset = dataset.map(parse_record)
164
  dataset = dataset.map(
165
166
      lambda image, label: (preprocess_image(image, is_training), label))

167
  dataset = dataset.prefetch(2 * batch_size)
168

Neal Wu's avatar
Neal Wu committed
169
170
  # We call repeat after shuffling, rather than before, to prevent separate
  # epochs from blending together.
171
  dataset = dataset.repeat(num_epochs)
172
173
174

  # Batch results by up to batch_size, and then fetch the tuple from the
  # iterator.
Neal Wu's avatar
Neal Wu committed
175
176
  dataset = dataset.batch(batch_size)
  iterator = dataset.make_one_shot_iterator()
177
178
179
180
181
  images, labels = iterator.get_next()

  return images, labels


182
def cifar10_model_fn(features, labels, mode, params):
183
184
185
186
  """Model function for CIFAR-10."""
  tf.summary.image('images', features, max_outputs=6)

  network = resnet_model.cifar10_resnet_v2_generator(
187
      params['resnet_size'], _NUM_CLASSES, params['data_format'])
188

189
  inputs = tf.reshape(features, [-1, _HEIGHT, _WIDTH, _DEPTH])
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
  logits = network(inputs, mode == tf.estimator.ModeKeys.TRAIN)

  predictions = {
      'classes': tf.argmax(logits, axis=1),
      'probabilities': tf.nn.softmax(logits, name='softmax_tensor')
  }

  if mode == tf.estimator.ModeKeys.PREDICT:
    return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)

  # Calculate loss, which includes softmax cross entropy and L2 regularization.
  cross_entropy = tf.losses.softmax_cross_entropy(
      logits=logits, onehot_labels=labels)

  # Create a tensor named cross_entropy for logging purposes.
  tf.identity(cross_entropy, name='cross_entropy')
  tf.summary.scalar('cross_entropy', cross_entropy)

  # Add weight decay to the loss.
  loss = cross_entropy + _WEIGHT_DECAY * tf.add_n(
      [tf.nn.l2_loss(v) for v in tf.trainable_variables()])

  if mode == tf.estimator.ModeKeys.TRAIN:
213
214
    # Scale the learning rate linearly with the batch size. When the batch size
    # is 128, the learning rate should be 0.1.
215
216
    initial_learning_rate = 0.1 * params['batch_size'] / 128
    batches_per_epoch = _NUM_IMAGES['train'] / params['batch_size']
217
218
219
    global_step = tf.train.get_or_create_global_step()

    # Multiply the learning rate by 0.1 at 100, 150, and 200 epochs.
220
221
    boundaries = [int(batches_per_epoch * epoch) for epoch in [100, 150, 200]]
    values = [initial_learning_rate * decay for decay in [1, 0.1, 0.01, 0.001]]
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
    learning_rate = tf.train.piecewise_constant(
        tf.cast(global_step, tf.int32), boundaries, values)

    # Create a tensor named learning_rate for logging purposes
    tf.identity(learning_rate, name='learning_rate')
    tf.summary.scalar('learning_rate', learning_rate)

    optimizer = tf.train.MomentumOptimizer(
        learning_rate=learning_rate,
        momentum=_MOMENTUM)

    # Batch norm requires update ops to be added as a dependency to the train_op
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
      train_op = optimizer.minimize(loss, global_step)
  else:
    train_op = None

240
  accuracy = tf.metrics.accuracy(
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
      tf.argmax(labels, axis=1), predictions['classes'])
  metrics = {'accuracy': accuracy}

  # Create a tensor named train_accuracy for logging purposes
  tf.identity(accuracy[1], name='train_accuracy')
  tf.summary.scalar('train_accuracy', accuracy[1])

  return tf.estimator.EstimatorSpec(
      mode=mode,
      predictions=predictions,
      loss=loss,
      train_op=train_op,
      eval_metric_ops=metrics)


def main(unused_argv):
  # Using the Winograd non-fused algorithms provides a small performance boost.
  os.environ['TF_ENABLE_WINOGRAD_NONFUSED'] = '1'

260
261
  # Set up a RunConfig to only save checkpoints once per training cycle.
  run_config = tf.estimator.RunConfig().replace(save_checkpoints_secs=1e9)
262
  cifar_classifier = tf.estimator.Estimator(
263
264
265
266
267
268
      model_fn=cifar10_model_fn, model_dir=FLAGS.model_dir, config=run_config,
      params={
          'resnet_size': FLAGS.resnet_size,
          'data_format': FLAGS.data_format,
          'batch_size': FLAGS.batch_size,
      })
269

270
  for _ in range(FLAGS.train_epochs // FLAGS.epochs_per_eval):
271
272
273
274
275
276
277
278
279
280
    tensors_to_log = {
        'learning_rate': 'learning_rate',
        'cross_entropy': 'cross_entropy',
        'train_accuracy': 'train_accuracy'
    }

    logging_hook = tf.train.LoggingTensorHook(
        tensors=tensors_to_log, every_n_iter=100)

    cifar_classifier.train(
281
        input_fn=lambda: input_fn(
282
            True, FLAGS.data_dir, FLAGS.batch_size, FLAGS.epochs_per_eval),
283
284
285
286
        hooks=[logging_hook])

    # Evaluate the model and print results
    eval_results = cifar_classifier.evaluate(
287
        input_fn=lambda: input_fn(False, FLAGS.data_dir, FLAGS.batch_size))
288
289
290
291
292
    print(eval_results)


if __name__ == '__main__':
  tf.logging.set_verbosity(tf.logging.INFO)
293
294
  FLAGS, unparsed = parser.parse_known_args()
  tf.app.run(argv=[sys.argv[0]] + unparsed)