imagenet_test.py 11.9 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# 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.
# ==============================================================================

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

import unittest

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

24
from official.resnet import imagenet_main
25
from official.utils.testing import integration
26
27
28

tf.logging.set_verbosity(tf.logging.ERROR)

29
_BATCH_SIZE = 32
30
31
32
33
34
_LABEL_CLASSES = 1001


class BaseTest(tf.test.TestCase):

35
36
37
38
39
  @classmethod
  def setUpClass(cls):  # pylint: disable=invalid-name
    super(BaseTest, cls).setUpClass()
    imagenet_main.define_imagenet_flags()

40
41
42
43
  def tearDown(self):
    super(BaseTest, self).tearDown()
    tf.gfile.DeleteRecursively(self.get_temp_dir())

44
  def _tensor_shapes_helper(self, resnet_size, version, dtype, with_gpu):
45
46
    """Checks the tensor shapes after each phase of the ResNet model."""
    def reshape(shape):
Karmel Allison's avatar
Karmel Allison committed
47
48
      """Returns the expected dimensions depending on if a GPU is being used."""

49
50
      # If a GPU is used for the test, the shape is returned (already in NCHW
      # form). When GPU is not used, the shape is converted to NHWC.
51
52
53
54
55
56
57
      if with_gpu:
        return shape
      return shape[0], shape[2], shape[3], shape[1]

    graph = tf.Graph()

    with graph.as_default(), self.test_session(
58
        graph=graph, use_gpu=with_gpu, force_gpu=with_gpu):
59
      model = imagenet_main.ImagenetModel(
60
          resnet_size=resnet_size,
61
          data_format='channels_first' if with_gpu else 'channels_last',
62
63
64
          version=version,
          dtype=dtype
      )
65
      inputs = tf.random_uniform([1, 224, 224, 3])
66
      output = model(inputs, training=True)
67

68
69
70
71
72
73
74
75
      initial_conv = graph.get_tensor_by_name('resnet_model/initial_conv:0')
      max_pool = graph.get_tensor_by_name('resnet_model/initial_max_pool:0')
      block_layer1 = graph.get_tensor_by_name('resnet_model/block_layer1:0')
      block_layer2 = graph.get_tensor_by_name('resnet_model/block_layer2:0')
      block_layer3 = graph.get_tensor_by_name('resnet_model/block_layer3:0')
      block_layer4 = graph.get_tensor_by_name('resnet_model/block_layer4:0')
      reduce_mean = graph.get_tensor_by_name('resnet_model/final_reduce_mean:0')
      dense = graph.get_tensor_by_name('resnet_model/final_dense:0')
76
77
78
79
80
81
82
83
84
85
86

      self.assertAllEqual(initial_conv.shape, reshape((1, 64, 112, 112)))
      self.assertAllEqual(max_pool.shape, reshape((1, 64, 56, 56)))

      # The number of channels after each block depends on whether we're
      # using the building_block or the bottleneck_block.
      if resnet_size < 50:
        self.assertAllEqual(block_layer1.shape, reshape((1, 64, 56, 56)))
        self.assertAllEqual(block_layer2.shape, reshape((1, 128, 28, 28)))
        self.assertAllEqual(block_layer3.shape, reshape((1, 256, 14, 14)))
        self.assertAllEqual(block_layer4.shape, reshape((1, 512, 7, 7)))
87
        self.assertAllEqual(reduce_mean.shape, reshape((1, 512, 1, 1)))
88
89
90
91
92
      else:
        self.assertAllEqual(block_layer1.shape, reshape((1, 256, 56, 56)))
        self.assertAllEqual(block_layer2.shape, reshape((1, 512, 28, 28)))
        self.assertAllEqual(block_layer3.shape, reshape((1, 1024, 14, 14)))
        self.assertAllEqual(block_layer4.shape, reshape((1, 2048, 7, 7)))
93
        self.assertAllEqual(reduce_mean.shape, reshape((1, 2048, 1, 1)))
94

95
96
      self.assertAllEqual(dense.shape, (1, _LABEL_CLASSES))
      self.assertAllEqual(output.shape, (1, _LABEL_CLASSES))
97

98
99
100
101
102
103
  def tensor_shapes_helper(self, resnet_size, version, with_gpu=False):
    self._tensor_shapes_helper(resnet_size=resnet_size, version=version,
                               dtype=tf.float32, with_gpu=with_gpu)
    self._tensor_shapes_helper(resnet_size=resnet_size, version=version,
                               dtype=tf.float16, with_gpu=with_gpu)

104
105
  def test_tensor_shapes_resnet_18_v1(self):
    self.tensor_shapes_helper(18, version=1)
106

107
108
  def test_tensor_shapes_resnet_18_v2(self):
    self.tensor_shapes_helper(18, version=2)
109

110
111
  def test_tensor_shapes_resnet_34_v1(self):
    self.tensor_shapes_helper(34, version=1)
112

113
114
  def test_tensor_shapes_resnet_34_v2(self):
    self.tensor_shapes_helper(34, version=2)
115

116
117
  def test_tensor_shapes_resnet_50_v1(self):
    self.tensor_shapes_helper(50, version=1)
118

119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
  def test_tensor_shapes_resnet_50_v2(self):
    self.tensor_shapes_helper(50, version=2)

  def test_tensor_shapes_resnet_101_v1(self):
    self.tensor_shapes_helper(101, version=1)

  def test_tensor_shapes_resnet_101_v2(self):
    self.tensor_shapes_helper(101, version=2)

  def test_tensor_shapes_resnet_152_v1(self):
    self.tensor_shapes_helper(152, version=1)

  def test_tensor_shapes_resnet_152_v2(self):
    self.tensor_shapes_helper(152, version=2)

  def test_tensor_shapes_resnet_200_v1(self):
    self.tensor_shapes_helper(200, version=1)

  def test_tensor_shapes_resnet_200_v2(self):
    self.tensor_shapes_helper(200, version=2)

  @unittest.skipUnless(tf.test.is_built_with_cuda(), 'requires GPU')
  def test_tensor_shapes_resnet_18_with_gpu_v1(self):
    self.tensor_shapes_helper(18, version=1, with_gpu=True)

  @unittest.skipUnless(tf.test.is_built_with_cuda(), 'requires GPU')
  def test_tensor_shapes_resnet_18_with_gpu_v2(self):
    self.tensor_shapes_helper(18, version=2, with_gpu=True)

  @unittest.skipUnless(tf.test.is_built_with_cuda(), 'requires GPU')
  def test_tensor_shapes_resnet_34_with_gpu_v1(self):
    self.tensor_shapes_helper(34, version=1, with_gpu=True)

  @unittest.skipUnless(tf.test.is_built_with_cuda(), 'requires GPU')
  def test_tensor_shapes_resnet_34_with_gpu_v2(self):
    self.tensor_shapes_helper(34, version=2, with_gpu=True)
155
156

  @unittest.skipUnless(tf.test.is_built_with_cuda(), 'requires GPU')
157
158
  def test_tensor_shapes_resnet_50_with_gpu_v1(self):
    self.tensor_shapes_helper(50, version=1, with_gpu=True)
159
160

  @unittest.skipUnless(tf.test.is_built_with_cuda(), 'requires GPU')
161
162
  def test_tensor_shapes_resnet_50_with_gpu_v2(self):
    self.tensor_shapes_helper(50, version=2, with_gpu=True)
163
164

  @unittest.skipUnless(tf.test.is_built_with_cuda(), 'requires GPU')
165
166
  def test_tensor_shapes_resnet_101_with_gpu_v1(self):
    self.tensor_shapes_helper(101, version=1, with_gpu=True)
167
168

  @unittest.skipUnless(tf.test.is_built_with_cuda(), 'requires GPU')
169
170
  def test_tensor_shapes_resnet_101_with_gpu_v2(self):
    self.tensor_shapes_helper(101, version=2, with_gpu=True)
171
172

  @unittest.skipUnless(tf.test.is_built_with_cuda(), 'requires GPU')
173
174
  def test_tensor_shapes_resnet_152_with_gpu_v1(self):
    self.tensor_shapes_helper(152, version=1, with_gpu=True)
175
176

  @unittest.skipUnless(tf.test.is_built_with_cuda(), 'requires GPU')
177
178
  def test_tensor_shapes_resnet_152_with_gpu_v2(self):
    self.tensor_shapes_helper(152, version=2, with_gpu=True)
179

180
181
182
183
184
185
186
187
  @unittest.skipUnless(tf.test.is_built_with_cuda(), 'requires GPU')
  def test_tensor_shapes_resnet_200_with_gpu_v1(self):
    self.tensor_shapes_helper(200, version=1, with_gpu=True)

  @unittest.skipUnless(tf.test.is_built_with_cuda(), 'requires GPU')
  def test_tensor_shapes_resnet_200_with_gpu_v2(self):
    self.tensor_shapes_helper(200, version=2, with_gpu=True)

188
  def resnet_model_fn_helper(self, mode, version, dtype):
189
    """Tests that the EstimatorSpec is given the appropriate arguments."""
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
    tf.train.create_global_step()

    input_fn = imagenet_main.get_synth_input_fn()
    dataset = input_fn(True, '', _BATCH_SIZE)
    iterator = dataset.make_one_shot_iterator()
    features, labels = iterator.get_next()
    spec = imagenet_main.imagenet_model_fn(
        features, labels, mode, {
            'dtype': dtype,
            'resnet_size': 50,
            'data_format': 'channels_last',
            'batch_size': _BATCH_SIZE,
            'version': version,
            'loss_scale': 128 if dtype == tf.float16 else 1,
        })

    predictions = spec.predictions
    self.assertAllEqual(predictions['probabilities'].shape,
                        (_BATCH_SIZE, _LABEL_CLASSES))
    self.assertEqual(predictions['probabilities'].dtype, tf.float32)
    self.assertAllEqual(predictions['classes'].shape, (_BATCH_SIZE,))
    self.assertEqual(predictions['classes'].dtype, tf.int64)

    if mode != tf.estimator.ModeKeys.PREDICT:
      loss = spec.loss
      self.assertAllEqual(loss.shape, ())
      self.assertEqual(loss.dtype, tf.float32)

    if mode == tf.estimator.ModeKeys.EVAL:
      eval_metric_ops = spec.eval_metric_ops
      self.assertAllEqual(eval_metric_ops['accuracy'][0].shape, ())
      self.assertAllEqual(eval_metric_ops['accuracy'][1].shape, ())
      self.assertEqual(eval_metric_ops['accuracy'][0].dtype, tf.float32)
      self.assertEqual(eval_metric_ops['accuracy'][1].dtype, tf.float32)
224

225
  def test_resnet_model_fn_train_mode_v1(self):
226
    self.resnet_model_fn_helper(tf.estimator.ModeKeys.TRAIN, version=1,
227
                                dtype=tf.float32)
228

229
  def test_resnet_model_fn_train_mode_v2(self):
230
    self.resnet_model_fn_helper(tf.estimator.ModeKeys.TRAIN, version=2,
231
                                dtype=tf.float32)
232
233

  def test_resnet_model_fn_eval_mode_v1(self):
234
235
    self.resnet_model_fn_helper(tf.estimator.ModeKeys.EVAL, version=1,
                                dtype=tf.float32)
236

237
  def test_resnet_model_fn_eval_mode_v2(self):
238
239
    self.resnet_model_fn_helper(tf.estimator.ModeKeys.EVAL, version=2,
                                dtype=tf.float32)
Karmel Allison's avatar
Karmel Allison committed
240

241
  def test_resnet_model_fn_predict_mode_v1(self):
242
243
    self.resnet_model_fn_helper(tf.estimator.ModeKeys.PREDICT, version=1,
                                dtype=tf.float32)
244

245
  def test_resnet_model_fn_predict_mode_v2(self):
246
247
    self.resnet_model_fn_helper(tf.estimator.ModeKeys.PREDICT, version=2,
                                dtype=tf.float32)
248

249
  def _test_imagenetmodel_shape(self, version):
Neal Wu's avatar
Neal Wu committed
250
251
252
    batch_size = 135
    num_classes = 246

253
254
255
256
257
258
259
260
261
262
263
    model = imagenet_main.ImagenetModel(
        50, data_format='channels_last', num_classes=num_classes,
        version=version)

    fake_input = tf.random_uniform([batch_size, 224, 224, 3])
    output = model(fake_input, training=True)

    self.assertAllEqual(output.shape, (batch_size, num_classes))

  def test_imagenetmodel_shape_v1(self):
    self._test_imagenetmodel_shape(version=1)
Neal Wu's avatar
Neal Wu committed
264

265
266
  def test_imagenetmodel_shape_v2(self):
    self._test_imagenetmodel_shape(version=2)
Neal Wu's avatar
Neal Wu committed
267

268
  def test_imagenet_end_to_end_synthetic_v1(self):
269
    integration.run_synthetic(
270
        main=imagenet_main.run_imagenet, tmp_root=self.get_temp_dir(),
271
272
        extra_flags=['-v', '1']
    )
273
274

  def test_imagenet_end_to_end_synthetic_v2(self):
275
    integration.run_synthetic(
276
        main=imagenet_main.run_imagenet, tmp_root=self.get_temp_dir(),
277
278
        extra_flags=['-v', '2']
    )
279
280

  def test_imagenet_end_to_end_synthetic_v1_tiny(self):
281
    integration.run_synthetic(
282
        main=imagenet_main.run_imagenet, tmp_root=self.get_temp_dir(),
283
284
        extra_flags=['-v', '1', '-rs', '18']
    )
285
286

  def test_imagenet_end_to_end_synthetic_v2_tiny(self):
287
    integration.run_synthetic(
288
        main=imagenet_main.run_imagenet, tmp_root=self.get_temp_dir(),
289
290
        extra_flags=['-v', '2', '-rs', '18']
    )
291
292

  def test_imagenet_end_to_end_synthetic_v1_huge(self):
293
    integration.run_synthetic(
294
        main=imagenet_main.run_imagenet, tmp_root=self.get_temp_dir(),
295
296
        extra_flags=['-v', '1', '-rs', '200']
    )
297
298

  def test_imagenet_end_to_end_synthetic_v2_huge(self):
299
    integration.run_synthetic(
300
        main=imagenet_main.run_imagenet, tmp_root=self.get_temp_dir(),
301
302
        extra_flags=['-v', '2', '-rs', '200']
    )
303

304

305
306
if __name__ == '__main__':
  tf.test.main()