imagenet_test.py 11.1 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, with_gpu=False):
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
53
      if with_gpu:
        return shape
      return shape[0], shape[2], shape[3], shape[1]

    graph = tf.Graph()

    with graph.as_default(), self.test_session(
        use_gpu=with_gpu, force_gpu=with_gpu):
54
55
      model = imagenet_main.ImagenetModel(
          resnet_size,
56
57
          data_format='channels_first' if with_gpu else 'channels_last',
          version=version)
58
      inputs = tf.random_uniform([1, 224, 224, 3])
59
      output = model(inputs, training=True)
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87

      initial_conv = graph.get_tensor_by_name('initial_conv:0')
      max_pool = graph.get_tensor_by_name('initial_max_pool:0')
      block_layer1 = graph.get_tensor_by_name('block_layer1:0')
      block_layer2 = graph.get_tensor_by_name('block_layer2:0')
      block_layer3 = graph.get_tensor_by_name('block_layer3:0')
      block_layer4 = graph.get_tensor_by_name('block_layer4:0')
      avg_pool = graph.get_tensor_by_name('final_avg_pool:0')
      dense = graph.get_tensor_by_name('final_dense:0')

      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)))
        self.assertAllEqual(avg_pool.shape, reshape((1, 512, 1, 1)))
      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)))
        self.assertAllEqual(avg_pool.shape, reshape((1, 2048, 1, 1)))

88
89
      self.assertAllEqual(dense.shape, (1, _LABEL_CLASSES))
      self.assertAllEqual(output.shape, (1, _LABEL_CLASSES))
90

91
92
  def test_tensor_shapes_resnet_18_v1(self):
    self.tensor_shapes_helper(18, version=1)
93

94
95
  def test_tensor_shapes_resnet_18_v2(self):
    self.tensor_shapes_helper(18, version=2)
96

97
98
  def test_tensor_shapes_resnet_34_v1(self):
    self.tensor_shapes_helper(34, version=1)
99

100
101
  def test_tensor_shapes_resnet_34_v2(self):
    self.tensor_shapes_helper(34, version=2)
102

103
104
  def test_tensor_shapes_resnet_50_v1(self):
    self.tensor_shapes_helper(50, version=1)
105

106
107
108
109
110
111
112
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
  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)
142
143

  @unittest.skipUnless(tf.test.is_built_with_cuda(), 'requires GPU')
144
145
  def test_tensor_shapes_resnet_50_with_gpu_v1(self):
    self.tensor_shapes_helper(50, version=1, with_gpu=True)
146
147

  @unittest.skipUnless(tf.test.is_built_with_cuda(), 'requires GPU')
148
149
  def test_tensor_shapes_resnet_50_with_gpu_v2(self):
    self.tensor_shapes_helper(50, version=2, with_gpu=True)
150
151

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

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

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

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

167
168
169
170
171
172
173
174
175
  @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)

  def resnet_model_fn_helper(self, mode, version, multi_gpu=False):
176
177
178
    """Tests that the EstimatorSpec is given the appropriate arguments."""
    tf.train.create_global_step()

179
180
181
182
    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()
183
    spec = imagenet_main.imagenet_model_fn(
184
185
186
187
        features, labels, mode, {
            'resnet_size': 50,
            'data_format': 'channels_last',
            'batch_size': _BATCH_SIZE,
188
            'version': version,
Karmel Allison's avatar
Karmel Allison committed
189
            'multi_gpu': multi_gpu,
190
        })
191
192
193

    predictions = spec.predictions
    self.assertAllEqual(predictions['probabilities'].shape,
194
                        (_BATCH_SIZE, _LABEL_CLASSES))
195
    self.assertEqual(predictions['probabilities'].dtype, tf.float32)
196
    self.assertAllEqual(predictions['classes'].shape, (_BATCH_SIZE,))
197
198
199
200
201
202
203
204
205
206
207
208
209
210
    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)

211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
  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)
227

228
229
  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
230

231
232
  def test_resnet_model_fn_predict_mode_v1(self):
    self.resnet_model_fn_helper(tf.estimator.ModeKeys.PREDICT, version=1)
233

234
235
  def test_resnet_model_fn_predict_mode_v2(self):
    self.resnet_model_fn_helper(tf.estimator.ModeKeys.PREDICT, version=2)
236

Neal Wu's avatar
Neal Wu committed
237
238
239
240
  def test_imagenetmodel_shape(self):
    batch_size = 135
    num_classes = 246

241
    for version in (1, 2):
Karmel Allison's avatar
Karmel Allison committed
242
243
244
      model = imagenet_main.ImagenetModel(
          50, data_format='channels_last', num_classes=num_classes,
          version=version)
245
246
      fake_input = tf.random_uniform([batch_size, 224, 224, 3])
      output = model(fake_input, training=True)
Neal Wu's avatar
Neal Wu committed
247

248
      self.assertAllEqual(output.shape, (batch_size, num_classes))
Neal Wu's avatar
Neal Wu committed
249

250
  def test_imagenet_end_to_end_synthetic_v1(self):
251
252
253
254
    integration.run_synthetic(
        main=imagenet_main.main, tmp_root=self.get_temp_dir(),
        extra_flags=['-v', '1']
    )
255
256

  def test_imagenet_end_to_end_synthetic_v2(self):
257
258
259
260
    integration.run_synthetic(
        main=imagenet_main.main, tmp_root=self.get_temp_dir(),
        extra_flags=['-v', '2']
    )
261
262

  def test_imagenet_end_to_end_synthetic_v1_tiny(self):
263
264
265
266
    integration.run_synthetic(
        main=imagenet_main.main, tmp_root=self.get_temp_dir(),
        extra_flags=['-v', '1', '-rs', '18']
    )
267
268

  def test_imagenet_end_to_end_synthetic_v2_tiny(self):
269
270
271
272
    integration.run_synthetic(
        main=imagenet_main.main, tmp_root=self.get_temp_dir(),
        extra_flags=['-v', '2', '-rs', '18']
    )
273
274

  def test_imagenet_end_to_end_synthetic_v1_huge(self):
275
276
277
278
    integration.run_synthetic(
        main=imagenet_main.main, tmp_root=self.get_temp_dir(),
        extra_flags=['-v', '1', '-rs', '200']
    )
279
280

  def test_imagenet_end_to_end_synthetic_v2_huge(self):
281
282
283
284
    integration.run_synthetic(
        main=imagenet_main.main, tmp_root=self.get_temp_dir(),
        extra_flags=['-v', '2', '-rs', '200']
    )
285
286
287

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