imagenet_main.py 6.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 ImageNet 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
import vgg_preprocessing

29
_DEFAULT_IMAGE_SIZE = 224
30
_NUM_CHANNELS = 3
31
_NUM_CLASSES = 1001
32

33
34
35
36
_NUM_IMAGES = {
    'train': 1281167,
    'validation': 50000,
}
37

38
_NUM_TRAIN_FILES = 1024
39
_SHUFFLE_BUFFER = 1500
40

41

42
43
44
###############################################################################
# Data processing
###############################################################################
45
def get_filenames(is_training, data_dir):
46
47
48
  """Return filenames for dataset."""
  if is_training:
    return [
49
        os.path.join(data_dir, 'train-%05d-of-01024' % i)
50
        for i in range(_NUM_TRAIN_FILES)]
51
52
  else:
    return [
53
        os.path.join(data_dir, 'validation-%05d-of-00128' % i)
Neal Wu's avatar
Neal Wu committed
54
        for i in range(128)]
55
56


57
def parse_record(raw_record, is_training):
58
  """Parse an ImageNet record from `value`."""
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
  keys_to_features = {
      'image/encoded':
          tf.FixedLenFeature((), tf.string, default_value=''),
      'image/format':
          tf.FixedLenFeature((), tf.string, default_value='jpeg'),
      'image/class/label':
          tf.FixedLenFeature([], dtype=tf.int64, default_value=-1),
      'image/class/text':
          tf.FixedLenFeature([], dtype=tf.string, default_value=''),
      'image/object/bbox/xmin':
          tf.VarLenFeature(dtype=tf.float32),
      'image/object/bbox/ymin':
          tf.VarLenFeature(dtype=tf.float32),
      'image/object/bbox/xmax':
          tf.VarLenFeature(dtype=tf.float32),
      'image/object/bbox/ymax':
          tf.VarLenFeature(dtype=tf.float32),
      'image/object/class/label':
          tf.VarLenFeature(dtype=tf.int64),
  }

80
  parsed = tf.parse_single_example(raw_record, keys_to_features)
81

82
83
84
  image = tf.image.decode_image(
      tf.reshape(parsed['image/encoded'], shape=[]),
      _NUM_CHANNELS)
85
86

  # Note that tf.image.convert_image_dtype scales the image data to [0, 1).
87
88
  image = tf.image.convert_image_dtype(image, dtype=tf.float32)

89
  image = vgg_preprocessing.preprocess_image(
90
      image=image,
91
92
      output_height=_DEFAULT_IMAGE_SIZE,
      output_width=_DEFAULT_IMAGE_SIZE,
93
94
95
96
97
98
      is_training=is_training)

  label = tf.cast(
      tf.reshape(parsed['image/class/label'], shape=[]),
      dtype=tf.int32)

99
  return image, tf.one_hot(label, _NUM_CLASSES)
100
101


102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
def input_fn(is_training, data_dir, batch_size, num_epochs=1,
             num_parallel_calls=1):
  """Input function which provides batches for train or eval.
  Args:
    is_training: A boolean denoting whether the input is for training.
    data_dir: The directory containing the input data.
    batch_size: The number of samples per batch.
    num_epochs: The number of epochs to repeat the dataset.
    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.

  Returns:
    A dataset that can be used for iteration.
  """
  filenames = get_filenames(is_training, data_dir)
  dataset = tf.data.Dataset.from_tensor_slices(filenames)
119

120
  if is_training:
121
122
    # Shuffle the input files
    dataset = dataset.shuffle(buffer_size=_NUM_TRAIN_FILES)
123

124
  # Convert to individual records
125
  dataset = dataset.flat_map(tf.data.TFRecordDataset)
126

127
128
  return resnet.process_record_dataset(dataset, is_training, batch_size,
      _SHUFFLE_BUFFER, parse_record, num_epochs, num_parallel_calls)
129
130


131
132
133
###############################################################################
# Running the model
###############################################################################
Karmel Allison's avatar
Karmel Allison committed
134
class ImagenetModel(resnet.Model):
135
136
137
138
139
140
  def __init__(self, resnet_size, data_format=None):
    """These are the parameters that work for Imagenet data.
    """

    # For bigger models, we want to use "bottleneck" layers
    if resnet_size < 50:
Karmel Allison's avatar
Karmel Allison committed
141
      block_fn = resnet.building_block
142
143
      final_size = 512
    else:
Karmel Allison's avatar
Karmel Allison committed
144
      block_fn = resnet.bottleneck_block
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
      final_size = 2048

    super(ImagenetModel, self).__init__(
        resnet_size=resnet_size,
        num_classes=_NUM_CLASSES,
        num_filters=64,
        kernel_size=7,
        conv_stride=2,
        first_pool_size=3,
        first_pool_stride=2,
        second_pool_size=7,
        second_pool_stride=1,
        block_fn=block_fn,
        block_sizes=_get_block_sizes(resnet_size),
        block_strides=[1, 2, 2, 2],
        final_size=final_size,
        data_format=data_format)


def _get_block_sizes(resnet_size):
  """The number of block layers used for the Resnet model varies according
  to the size of the model. This helper grabs the layer set we want, throwing
  an error if a non-standard size has been selected.
  """
  choices = {
      18: [2, 2, 2, 2],
      34: [3, 4, 6, 3],
      50: [3, 4, 6, 3],
      101: [3, 4, 23, 3],
      152: [3, 8, 36, 3],
      200: [3, 24, 36, 3]
176
177
  }

178
179
180
181
182
183
184
  try:
    return choices[resnet_size]
  except KeyError:
    err = ('Could not find layers for selected Resnet size.\n'
           'Size received: {}; sizes allowed: {}.'.format(
               resnet_size, choices.keys()))
    raise ValueError(err)
185
186


187
188
def imagenet_model_fn(features, labels, mode, params):
  """Our model_fn for ResNet to be used with our Estimator."""
Karmel Allison's avatar
Karmel Allison committed
189
  learning_rate_fn = resnet.learning_rate_with_decay(
190
191
192
      batch_size=params['batch_size'], batch_denom=256,
      num_images=_NUM_IMAGES['train'], boundary_epochs=[30, 60, 80, 90],
      decay_rates=[1, 0.1, 0.01, 0.001, 1e-4])
193

Karmel Allison's avatar
Karmel Allison committed
194
195
196
197
198
199
200
  return resnet.resnet_model_fn(features, labels, mode, ImagenetModel,
                                resnet_size=params['resnet_size'],
                                weight_decay=1e-4,
                                learning_rate_fn=learning_rate_fn,
                                momentum=0.9,
                                data_format=params['data_format'],
                                loss_filter_fn=None)
201
202
203


def main(unused_argv):
Karmel Allison's avatar
Karmel Allison committed
204
  resnet.resnet_main(FLAGS, imagenet_model_fn, input_fn)
205
206
207
208


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

Karmel Allison's avatar
Karmel Allison committed
210
  parser = resnet.ResnetArgParser(
211
      resnet_size_choices=[18, 34, 50, 101, 152, 200])
212
213
  FLAGS, unparsed = parser.parse_known_args()
  tf.app.run(argv=[sys.argv[0]] + unparsed)