resnet_imagenet_test.py 8.37 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
# 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.
# ==============================================================================
"""Test the keras ResNet model with ImageNet data."""

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

from absl.testing import parameterized
import tensorflow as tf

from tensorflow.python.eager import context
from official.benchmark.models import resnet_imagenet_main
Fan Yang's avatar
Fan Yang committed
26
from official.legacy.image_classification.resnet import imagenet_preprocessing
27
28
from official.utils.testing import integration

29

30
31
32
33
@parameterized.parameters(
    "resnet",
    # "resnet_polynomial_decay",  b/151854314
    "mobilenet",
34
35
    # "mobilenet_polynomial_decay",  b/151854314
    "mobilenet_selective_clustering",
36
37
38
39
)
class KerasImagenetTest(tf.test.TestCase):
  """Unit tests for Keras Models with ImageNet."""
  _default_flags_dict = [
Hongkun Yu's avatar
Hongkun Yu committed
40
41
42
43
44
45
46
47
      "-batch_size",
      "4",
      "-train_steps",
      "1",
      "-use_synthetic_data",
      "true",
      "-data_format",
      "channels_last",
48
49
50
  ]
  _extra_flags_dict = {
      "resnet": [
Hongkun Yu's avatar
Hongkun Yu committed
51
52
53
54
          "-model",
          "resnet50_v1.5",
          "-optimizer",
          "resnet50_default",
55
56
      ],
      "resnet_polynomial_decay": [
Hongkun Yu's avatar
Hongkun Yu committed
57
58
59
60
61
62
          "-model",
          "resnet50_v1.5",
          "-optimizer",
          "resnet50_default",
          "-pruning_method",
          "polynomial_decay",
63
64
      ],
      "mobilenet": [
Hongkun Yu's avatar
Hongkun Yu committed
65
66
67
68
          "-model",
          "mobilenet",
          "-optimizer",
          "mobilenet_default",
69
70
      ],
      "mobilenet_polynomial_decay": [
Hongkun Yu's avatar
Hongkun Yu committed
71
72
73
74
75
76
          "-model",
          "mobilenet",
          "-optimizer",
          "mobilenet_default",
          "-pruning_method",
          "polynomial_decay",
77
      ],
78
79
80
81
82
      "mobilenet_selective_clustering": [
          "-model", "mobilenet_pretrained",
          "-optimizer", "mobilenet_fine_tune",
          "-clustering_method", "selective_clustering",
      ]
83
84
85
86
87
88
89
90
91
92
93
  }
  _tempdir = None

  @classmethod
  def setUpClass(cls):  # pylint: disable=invalid-name
    super(KerasImagenetTest, cls).setUpClass()
    resnet_imagenet_main.define_imagenet_keras_flags()

  def setUp(self):
    super(KerasImagenetTest, self).setUp()
    imagenet_preprocessing.NUM_IMAGES["validation"] = 4
94
    self.policy = tf.keras.mixed_precision.global_policy()
95
96
97
98

  def tearDown(self):
    super(KerasImagenetTest, self).tearDown()
    tf.io.gfile.rmtree(self.get_temp_dir())
99
    tf.keras.mixed_precision.set_global_policy(self.policy)
100
101
102
103
104
105
106
107

  def get_extra_flags_dict(self, flags_key):
    return self._extra_flags_dict[flags_key] + self._default_flags_dict

  def test_end_to_end_no_dist_strat(self, flags_key):
    """Test Keras model with 1 GPU, no distribution strategy."""

    extra_flags = [
Hongkun Yu's avatar
Hongkun Yu committed
108
109
        "-distribution_strategy",
        "off",
110
111
112
113
114
115
    ]
    extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)

    integration.run_synthetic(
        main=resnet_imagenet_main.run,
        tmp_root=self.get_temp_dir(),
Hongkun Yu's avatar
Hongkun Yu committed
116
        extra_flags=extra_flags)
117
118
119
120

  def test_end_to_end_graph_no_dist_strat(self, flags_key):
    """Test Keras model in legacy graph mode with 1 GPU, no dist strat."""
    extra_flags = [
Hongkun Yu's avatar
Hongkun Yu committed
121
122
123
124
        "-enable_eager",
        "false",
        "-distribution_strategy",
        "off",
125
126
127
128
129
130
    ]
    extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)

    integration.run_synthetic(
        main=resnet_imagenet_main.run,
        tmp_root=self.get_temp_dir(),
Hongkun Yu's avatar
Hongkun Yu committed
131
        extra_flags=extra_flags)
132
133
134
135
136
137

  def test_end_to_end_1_gpu(self, flags_key):
    """Test Keras model with 1 GPU."""

    if context.num_gpus() < 1:
      self.skipTest(
Hongkun Yu's avatar
Hongkun Yu committed
138
139
          "{} GPUs are not available for this test. {} GPUs are available"
          .format(1, context.num_gpus()))
140
141

    extra_flags = [
Hongkun Yu's avatar
Hongkun Yu committed
142
143
144
145
146
147
        "-num_gpus",
        "1",
        "-distribution_strategy",
        "mirrored",
        "-enable_checkpoint_and_export",
        "1",
148
149
150
151
152
153
    ]
    extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)

    integration.run_synthetic(
        main=resnet_imagenet_main.run,
        tmp_root=self.get_temp_dir(),
Hongkun Yu's avatar
Hongkun Yu committed
154
        extra_flags=extra_flags)
155
156
157
158
159
160
161
162
163
164

  def test_end_to_end_1_gpu_fp16(self, flags_key):
    """Test Keras model with 1 GPU and fp16."""

    if context.num_gpus() < 1:
      self.skipTest(
          "{} GPUs are not available for this test. {} GPUs are available"
          .format(1, context.num_gpus()))

    extra_flags = [
Hongkun Yu's avatar
Hongkun Yu committed
165
166
167
168
169
170
        "-num_gpus",
        "1",
        "-dtype",
        "fp16",
        "-distribution_strategy",
        "mirrored",
171
172
173
174
    ]
    extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)

    if "polynomial_decay" in extra_flags:
175
176
177
178
      self.skipTest("Pruning with fp16 is currently not supported.")

    if "selective_clustering" in extra_flags:
      self.skipTest("Clustering with fp16 is currently not supported.")
179
180
181
182

    integration.run_synthetic(
        main=resnet_imagenet_main.run,
        tmp_root=self.get_temp_dir(),
Hongkun Yu's avatar
Hongkun Yu committed
183
        extra_flags=extra_flags)
184
185
186
187
188
189

  def test_end_to_end_2_gpu(self, flags_key):
    """Test Keras model with 2 GPUs."""

    if context.num_gpus() < 2:
      self.skipTest(
Hongkun Yu's avatar
Hongkun Yu committed
190
191
          "{} GPUs are not available for this test. {} GPUs are available"
          .format(2, context.num_gpus()))
192
193

    extra_flags = [
Hongkun Yu's avatar
Hongkun Yu committed
194
195
196
197
        "-num_gpus",
        "2",
        "-distribution_strategy",
        "mirrored",
198
199
200
201
202
203
    ]
    extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)

    integration.run_synthetic(
        main=resnet_imagenet_main.run,
        tmp_root=self.get_temp_dir(),
Hongkun Yu's avatar
Hongkun Yu committed
204
        extra_flags=extra_flags)
205
206
207
208
209
210

  def test_end_to_end_xla_2_gpu(self, flags_key):
    """Test Keras model with XLA and 2 GPUs."""

    if context.num_gpus() < 2:
      self.skipTest(
Hongkun Yu's avatar
Hongkun Yu committed
211
212
          "{} GPUs are not available for this test. {} GPUs are available"
          .format(2, context.num_gpus()))
213
214

    extra_flags = [
Hongkun Yu's avatar
Hongkun Yu committed
215
216
217
218
219
220
        "-num_gpus",
        "2",
        "-enable_xla",
        "true",
        "-distribution_strategy",
        "mirrored",
221
222
223
224
225
226
    ]
    extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)

    integration.run_synthetic(
        main=resnet_imagenet_main.run,
        tmp_root=self.get_temp_dir(),
Hongkun Yu's avatar
Hongkun Yu committed
227
        extra_flags=extra_flags)
228
229
230
231
232
233

  def test_end_to_end_2_gpu_fp16(self, flags_key):
    """Test Keras model with 2 GPUs and fp16."""

    if context.num_gpus() < 2:
      self.skipTest(
Hongkun Yu's avatar
Hongkun Yu committed
234
235
          "{} GPUs are not available for this test. {} GPUs are available"
          .format(2, context.num_gpus()))
236
237

    extra_flags = [
Hongkun Yu's avatar
Hongkun Yu committed
238
239
240
241
242
243
        "-num_gpus",
        "2",
        "-dtype",
        "fp16",
        "-distribution_strategy",
        "mirrored",
244
245
246
247
    ]
    extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)

    if "polynomial_decay" in extra_flags:
248
249
250
251
      self.skipTest("Pruning with fp16 is currently not supported.")

    if "selective_clustering" in extra_flags:
      self.skipTest("Clustering with fp16 is currently not supported.")
252
253
254
255

    integration.run_synthetic(
        main=resnet_imagenet_main.run,
        tmp_root=self.get_temp_dir(),
Hongkun Yu's avatar
Hongkun Yu committed
256
        extra_flags=extra_flags)
257
258
259
260
261

  def test_end_to_end_xla_2_gpu_fp16(self, flags_key):
    """Test Keras model with XLA, 2 GPUs and fp16."""
    if context.num_gpus() < 2:
      self.skipTest(
Hongkun Yu's avatar
Hongkun Yu committed
262
263
          "{} GPUs are not available for this test. {} GPUs are available"
          .format(2, context.num_gpus()))
264
265

    extra_flags = [
Hongkun Yu's avatar
Hongkun Yu committed
266
267
268
269
270
271
272
273
        "-num_gpus",
        "2",
        "-dtype",
        "fp16",
        "-enable_xla",
        "true",
        "-distribution_strategy",
        "mirrored",
274
275
276
277
    ]
    extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)

    if "polynomial_decay" in extra_flags:
278
279
280
281
      self.skipTest("Pruning with fp16 is currently not supported.")

    if "selective_clustering" in extra_flags:
      self.skipTest("Clustering with fp16 is currently not supported.")
282
283
284
285

    integration.run_synthetic(
        main=resnet_imagenet_main.run,
        tmp_root=self.get_temp_dir(),
Hongkun Yu's avatar
Hongkun Yu committed
286
        extra_flags=extra_flags)
287
288
289
290


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