imagenet_test.py 12.8 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
  def tearDown(self):
    super(BaseTest, self).tearDown()
    tf.gfile.DeleteRecursively(self.get_temp_dir())

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

44
45
      # 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.
46
47
48
49
50
51
52
      if with_gpu:
        return shape
      return shape[0], shape[2], shape[3], shape[1]

    graph = tf.Graph()

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

63
64
65
66
67
68
69
70
      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')
71
72
73
74
75
76
77
78
79
80
81

      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)))
82
        self.assertAllEqual(reduce_mean.shape, reshape((1, 512, 1, 1)))
83
84
85
86
87
      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)))
88
        self.assertAllEqual(reduce_mean.shape, reshape((1, 2048, 1, 1)))
89

90
91
      self.assertAllEqual(dense.shape, (1, _LABEL_CLASSES))
      self.assertAllEqual(output.shape, (1, _LABEL_CLASSES))
92

93
94
95
96
97
98
  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)

99
100
  def test_tensor_shapes_resnet_18_v1(self):
    self.tensor_shapes_helper(18, version=1)
101

102
103
  def test_tensor_shapes_resnet_18_v2(self):
    self.tensor_shapes_helper(18, version=2)
104

105
106
  def test_tensor_shapes_resnet_34_v1(self):
    self.tensor_shapes_helper(34, version=1)
107

108
109
  def test_tensor_shapes_resnet_34_v2(self):
    self.tensor_shapes_helper(34, version=2)
110

111
112
  def test_tensor_shapes_resnet_50_v1(self):
    self.tensor_shapes_helper(50, version=1)
113

114
115
116
117
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
  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)
150
151

  @unittest.skipUnless(tf.test.is_built_with_cuda(), 'requires GPU')
152
153
  def test_tensor_shapes_resnet_50_with_gpu_v1(self):
    self.tensor_shapes_helper(50, version=1, with_gpu=True)
154
155

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

  @unittest.skipUnless(tf.test.is_built_with_cuda(), 'requires GPU')
160
161
  def test_tensor_shapes_resnet_101_with_gpu_v1(self):
    self.tensor_shapes_helper(101, version=1, with_gpu=True)
162
163

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

  @unittest.skipUnless(tf.test.is_built_with_cuda(), 'requires GPU')
168
169
  def test_tensor_shapes_resnet_152_with_gpu_v1(self):
    self.tensor_shapes_helper(152, version=1, with_gpu=True)
170
171

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

175
176
177
178
179
180
181
182
  @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)

183
  def _resnet_model_fn_helper(self, mode, version, dtype, multi_gpu):
184
    """Tests that the EstimatorSpec is given the appropriate arguments."""
185
186
187
188
189
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
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
    with tf.Graph().as_default() as g:
      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,
              'multi_gpu': multi_gpu,
          })

      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)

        tensors_to_check = ('initial_conv:0', 'initial_max_pool:0',
                            'block_layer1:0', 'block_layer2:0',
                            'block_layer3:0', 'block_layer4:0',
                            'final_reduce_mean:0', 'final_dense:0')

        for tensor_name in tensors_to_check:
          tensor = g.get_tensor_by_name('resnet_model/' + tensor_name)
          self.assertEqual(tensor.dtype, dtype,
                           'Tensor {} has dtype {}, while dtype {} was '
                           'expected'.format(tensor, tensor.dtype,
                                             dtype))

  def resnet_model_fn_helper(self, mode, version, multi_gpu=False):
    self._resnet_model_fn_helper(mode=mode, version=version, dtype=tf.float32,
                                 multi_gpu=multi_gpu)
    self._resnet_model_fn_helper(mode=mode, version=version, dtype=tf.float16,
                                 multi_gpu=multi_gpu)
239

240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
  def test_resnet_model_fn_train_mode_v1(self):
    self.resnet_model_fn_helper(tf.estimator.ModeKeys.TRAIN, version=1)

  def test_resnet_model_fn_train_mode_v2(self):
    self.resnet_model_fn_helper(tf.estimator.ModeKeys.TRAIN, version=2)

  def test_resnet_model_fn_train_mode_multi_gpu_v1(self):
    self.resnet_model_fn_helper(tf.estimator.ModeKeys.TRAIN, version=1,
                                multi_gpu=True)

  def test_resnet_model_fn_train_mode_multi_gpu_v2(self):
    self.resnet_model_fn_helper(tf.estimator.ModeKeys.TRAIN, version=2,
                                multi_gpu=True)

  def test_resnet_model_fn_eval_mode_v1(self):
    self.resnet_model_fn_helper(tf.estimator.ModeKeys.EVAL, version=1)
256

257
258
  def test_resnet_model_fn_eval_mode_v2(self):
    self.resnet_model_fn_helper(tf.estimator.ModeKeys.EVAL, version=2)
Karmel Allison's avatar
Karmel Allison committed
259

260
261
  def test_resnet_model_fn_predict_mode_v1(self):
    self.resnet_model_fn_helper(tf.estimator.ModeKeys.PREDICT, version=1)
262

263
264
  def test_resnet_model_fn_predict_mode_v2(self):
    self.resnet_model_fn_helper(tf.estimator.ModeKeys.PREDICT, version=2)
265

266
  def _test_imagenetmodel_shape(self, version):
Neal Wu's avatar
Neal Wu committed
267
268
269
    batch_size = 135
    num_classes = 246

270
271
272
273
274
275
276
277
278
279
280
    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
281

282
283
  def test_imagenetmodel_shape_v2(self):
    self._test_imagenetmodel_shape(version=2)
Neal Wu's avatar
Neal Wu committed
284

285
  def test_imagenet_end_to_end_synthetic_v1(self):
286
287
288
289
    integration.run_synthetic(
        main=imagenet_main.main, tmp_root=self.get_temp_dir(),
        extra_flags=['-v', '1']
    )
290
291

  def test_imagenet_end_to_end_synthetic_v2(self):
292
293
294
295
    integration.run_synthetic(
        main=imagenet_main.main, tmp_root=self.get_temp_dir(),
        extra_flags=['-v', '2']
    )
296
297

  def test_imagenet_end_to_end_synthetic_v1_tiny(self):
298
299
300
301
    integration.run_synthetic(
        main=imagenet_main.main, tmp_root=self.get_temp_dir(),
        extra_flags=['-v', '1', '-rs', '18']
    )
302
303

  def test_imagenet_end_to_end_synthetic_v2_tiny(self):
304
305
306
307
    integration.run_synthetic(
        main=imagenet_main.main, tmp_root=self.get_temp_dir(),
        extra_flags=['-v', '2', '-rs', '18']
    )
308
309

  def test_imagenet_end_to_end_synthetic_v1_huge(self):
310
311
312
313
    integration.run_synthetic(
        main=imagenet_main.main, tmp_root=self.get_temp_dir(),
        extra_flags=['-v', '1', '-rs', '200']
    )
314
315

  def test_imagenet_end_to_end_synthetic_v2_huge(self):
316
317
318
319
    integration.run_synthetic(
        main=imagenet_main.main, tmp_root=self.get_temp_dir(),
        extra_flags=['-v', '2', '-rs', '200']
    )
320
321
322

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