resnet_cifar_test.py 5.46 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# Copyright 2018 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 Cifar data."""

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

21
22
import tempfile

23
24
25
26
import tensorflow as tf

from tensorflow.python.eager import context
from tensorflow.python.platform import googletest
27
from official.benchmark.models import cifar_preprocessing
28
from official.benchmark.models import resnet_cifar_main
29
from official.utils.testing import integration
30
31
32
33
34
35


class KerasCifarTest(googletest.TestCase):
  """Unit tests for Keras ResNet with Cifar."""

  _extra_flags = [
Hongkun Yu's avatar
Hongkun Yu committed
36
      "-batch_size", "4", "-train_steps", "1", "-use_synthetic_data", "true"
37
38
39
40
41
  ]
  _tempdir = None

  def get_temp_dir(self):
    if not self._tempdir:
42
      self._tempdir = tempfile.mkdtemp(dir=googletest.GetTempDir())
43
44
45
46
47
    return self._tempdir

  @classmethod
  def setUpClass(cls):  # pylint: disable=invalid-name
    super(KerasCifarTest, cls).setUpClass()
48
    resnet_cifar_main.define_cifar_flags()
49
50
51

  def setUp(self):
    super(KerasCifarTest, self).setUp()
52
    cifar_preprocessing.NUM_IMAGES["validation"] = 4
53
54
55
56
57
58
59
60
61

  def tearDown(self):
    super(KerasCifarTest, self).tearDown()
    tf.io.gfile.rmtree(self.get_temp_dir())

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

    extra_flags = [
Hongkun Yu's avatar
Hongkun Yu committed
62
63
64
65
66
67
        "-distribution_strategy",
        "off",
        "-model_dir",
        "keras_cifar_no_dist_strat",
        "-data_format",
        "channels_last",
68
69
70
71
    ]
    extra_flags = extra_flags + self._extra_flags

    integration.run_synthetic(
72
        main=resnet_cifar_main.run,
73
        tmp_root=self.get_temp_dir(),
Hongkun Yu's avatar
Hongkun Yu committed
74
        extra_flags=extra_flags)
75
76
77
78

  def test_end_to_end_graph_no_dist_strat(self):
    """Test Keras model in legacy graph mode with 1 GPU, no dist strat."""
    extra_flags = [
Hongkun Yu's avatar
Hongkun Yu committed
79
80
81
82
83
84
85
86
        "-enable_eager",
        "false",
        "-distribution_strategy",
        "off",
        "-model_dir",
        "keras_cifar_graph_no_dist_strat",
        "-data_format",
        "channels_last",
87
88
89
90
    ]
    extra_flags = extra_flags + self._extra_flags

    integration.run_synthetic(
91
        main=resnet_cifar_main.run,
92
        tmp_root=self.get_temp_dir(),
Hongkun Yu's avatar
Hongkun Yu committed
93
        extra_flags=extra_flags)
94
95
96
97
98
99

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

    if context.num_gpus() < 1:
      self.skipTest(
Hongkun Yu's avatar
Hongkun Yu committed
100
101
          "{} GPUs are not available for this test. {} GPUs are available"
          .format(1, context.num_gpus()))
102
103

    extra_flags = [
Hongkun Yu's avatar
Hongkun Yu committed
104
105
106
107
108
109
110
111
        "-num_gpus",
        "1",
        "-distribution_strategy",
        "mirrored",
        "-model_dir",
        "keras_cifar_1_gpu",
        "-data_format",
        "channels_last",
112
113
114
115
    ]
    extra_flags = extra_flags + self._extra_flags

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

  def test_end_to_end_graph_1_gpu(self):
    """Test Keras model in legacy graph mode with 1 GPU."""
    if context.num_gpus() < 1:
      self.skipTest(
Hongkun Yu's avatar
Hongkun Yu committed
124
125
          "{} GPUs are not available for this test. {} GPUs are available"
          .format(1, context.num_gpus()))
126
127

    extra_flags = [
Hongkun Yu's avatar
Hongkun Yu committed
128
129
        "-num_gpus",
        "1",
130
        "-noenable_eager",
Hongkun Yu's avatar
Hongkun Yu committed
131
132
133
134
135
136
        "-distribution_strategy",
        "mirrored",
        "-model_dir",
        "keras_cifar_graph_1_gpu",
        "-data_format",
        "channels_last",
137
138
139
140
    ]
    extra_flags = extra_flags + self._extra_flags

    integration.run_synthetic(
141
        main=resnet_cifar_main.run,
142
        tmp_root=self.get_temp_dir(),
Hongkun Yu's avatar
Hongkun Yu committed
143
        extra_flags=extra_flags)
144
145
146
147
148
149

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

    if context.num_gpus() < 2:
      self.skipTest(
Hongkun Yu's avatar
Hongkun Yu committed
150
151
          "{} GPUs are not available for this test. {} GPUs are available"
          .format(2, context.num_gpus()))
152
153

    extra_flags = [
Hongkun Yu's avatar
Hongkun Yu committed
154
155
156
157
158
159
        "-num_gpus",
        "2",
        "-distribution_strategy",
        "mirrored",
        "-model_dir",
        "keras_cifar_2_gpu",
160
161
162
163
    ]
    extra_flags = extra_flags + self._extra_flags

    integration.run_synthetic(
164
        main=resnet_cifar_main.run,
165
        tmp_root=self.get_temp_dir(),
Hongkun Yu's avatar
Hongkun Yu committed
166
        extra_flags=extra_flags)
167
168
169
170
171

  def test_end_to_end_graph_2_gpu(self):
    """Test Keras model in legacy graph mode with 2 GPUs."""
    if context.num_gpus() < 2:
      self.skipTest(
Hongkun Yu's avatar
Hongkun Yu committed
172
173
          "{} GPUs are not available for this test. {} GPUs are available"
          .format(2, context.num_gpus()))
174
175

    extra_flags = [
Hongkun Yu's avatar
Hongkun Yu committed
176
177
178
179
180
181
182
183
        "-num_gpus",
        "2",
        "-enable_eager",
        "false",
        "-distribution_strategy",
        "mirrored",
        "-model_dir",
        "keras_cifar_graph_2_gpu",
184
185
186
187
    ]
    extra_flags = extra_flags + self._extra_flags

    integration.run_synthetic(
188
        main=resnet_cifar_main.run,
189
        tmp_root=self.get_temp_dir(),
Hongkun Yu's avatar
Hongkun Yu committed
190
        extra_flags=extra_flags)
191
192
193
194


if __name__ == "__main__":
  googletest.main()