flags_test.py 5.21 KB
Newer Older
Hongkun Yu's avatar
Hongkun Yu committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
# Copyright 2021 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.

Taylor Robie's avatar
Taylor Robie committed
15
16
import unittest

17
18
from absl import flags
import tensorflow as tf
Taylor Robie's avatar
Taylor Robie committed
19

20
from official.utils.flags import core as flags_core  # pylint: disable=g-bad-import-order
Taylor Robie's avatar
Taylor Robie committed
21
22


23
def define_flags():
Hongkun Yu's avatar
Hongkun Yu committed
24
25
26
27
28
29
30
  flags_core.define_base(
      clean=True,
      num_gpu=False,
      stop_threshold=True,
      hooks=True,
      train_epochs=True,
      epochs_between_evals=True)
31
  flags_core.define_performance(
Hongkun Yu's avatar
Hongkun Yu committed
32
33
34
35
36
37
      num_parallel_calls=True,
      inter_op=True,
      intra_op=True,
      dynamic_loss_scale=True,
      loss_scale=True,
      synthetic_data=True,
38
      dtype=True)
39
40
  flags_core.define_image()
  flags_core.define_benchmark()
Taylor Robie's avatar
Taylor Robie committed
41
42
43
44


class BaseTester(unittest.TestCase):

45
46
47
48
49
  @classmethod
  def setUpClass(cls):
    super(BaseTester, cls).setUpClass()
    define_flags()

Taylor Robie's avatar
Taylor Robie committed
50
  def test_default_setting(self):
Hongkun Yu's avatar
Hongkun Yu committed
51
    """Test to ensure fields exist and defaults can be set."""
Taylor Robie's avatar
Taylor Robie committed
52
53
54
55
56

    defaults = dict(
        data_dir="dfgasf",
        model_dir="dfsdkjgbs",
        train_epochs=534,
57
        epochs_between_evals=15,
Taylor Robie's avatar
Taylor Robie committed
58
        batch_size=256,
59
        hooks=["LoggingTensorHook"],
Taylor Robie's avatar
Taylor Robie committed
60
61
        num_parallel_calls=18,
        inter_op_parallelism_threads=5,
62
        intra_op_parallelism_threads=10,
Hongkun Yu's avatar
Hongkun Yu committed
63
        data_format="channels_first")
Taylor Robie's avatar
Taylor Robie committed
64

65
66
    flags_core.set_defaults(**defaults)
    flags_core.parse_flags()
Taylor Robie's avatar
Taylor Robie committed
67
68

    for key, value in defaults.items():
69
      assert flags.FLAGS.get_flag_value(name=key, default=None) == value
Taylor Robie's avatar
Taylor Robie committed
70

71
72
73
  def test_benchmark_setting(self):
    defaults = dict(
        hooks=["LoggingMetricHook"],
74
75
        benchmark_log_dir="/tmp/12345",
        gcp_project="project_abc",
76
77
    )

78
79
    flags_core.set_defaults(**defaults)
    flags_core.parse_flags()
80
81

    for key, value in defaults.items():
82
      assert flags.FLAGS.get_flag_value(name=key, default=None) == value
83

Taylor Robie's avatar
Taylor Robie committed
84
  def test_booleans(self):
Hongkun Yu's avatar
Hongkun Yu committed
85
    """Test to ensure boolean flags trigger as expected."""
Taylor Robie's avatar
Taylor Robie committed
86

87
    flags_core.parse_flags([__file__, "--use_synthetic_data"])
Taylor Robie's avatar
Taylor Robie committed
88

89
    assert flags.FLAGS.use_synthetic_data
Taylor Robie's avatar
Taylor Robie committed
90

91
  def test_parse_dtype_info(self):
92
93
    flags_core.parse_flags([__file__, "--dtype", "fp16"])
    self.assertEqual(flags_core.get_tf_dtype(flags.FLAGS), tf.float16)
Hongkun Yu's avatar
Hongkun Yu committed
94
95
    self.assertEqual(
        flags_core.get_loss_scale(flags.FLAGS, default_for_fp16=2), 2)
96

Hongkun Yu's avatar
Hongkun Yu committed
97
98
99
    flags_core.parse_flags([__file__, "--dtype", "fp16", "--loss_scale", "5"])
    self.assertEqual(
        flags_core.get_loss_scale(flags.FLAGS, default_for_fp16=2), 5)
100
101
102

    flags_core.parse_flags(
        [__file__, "--dtype", "fp16", "--loss_scale", "dynamic"])
Hongkun Yu's avatar
Hongkun Yu committed
103
104
    self.assertEqual(
        flags_core.get_loss_scale(flags.FLAGS, default_for_fp16=2), "dynamic")
105
106
107

    flags_core.parse_flags([__file__, "--dtype", "fp32"])
    self.assertEqual(flags_core.get_tf_dtype(flags.FLAGS), tf.float32)
Hongkun Yu's avatar
Hongkun Yu committed
108
109
    self.assertEqual(
        flags_core.get_loss_scale(flags.FLAGS, default_for_fp16=2), 1)
110
111

    flags_core.parse_flags([__file__, "--dtype", "fp32", "--loss_scale", "5"])
Hongkun Yu's avatar
Hongkun Yu committed
112
113
    self.assertEqual(
        flags_core.get_loss_scale(flags.FLAGS, default_for_fp16=2), 5)
114

115
    with self.assertRaises(SystemExit):
116
      flags_core.parse_flags([__file__, "--dtype", "int8"])
117

118
    with self.assertRaises(SystemExit):
Hongkun Yu's avatar
Hongkun Yu committed
119
120
      flags_core.parse_flags(
          [__file__, "--dtype", "fp16", "--loss_scale", "abc"])
121

122
123
124
125
126
127
  def test_get_nondefault_flags_as_str(self):
    defaults = dict(
        clean=True,
        data_dir="abc",
        hooks=["LoggingTensorHook"],
        stop_threshold=1.5,
Hongkun Yu's avatar
Hongkun Yu committed
128
        use_synthetic_data=False)
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
155
156
157
158
159
    flags_core.set_defaults(**defaults)
    flags_core.parse_flags()

    expected_flags = ""
    self.assertEqual(flags_core.get_nondefault_flags_as_str(), expected_flags)

    flags.FLAGS.clean = False
    expected_flags += "--noclean"
    self.assertEqual(flags_core.get_nondefault_flags_as_str(), expected_flags)

    flags.FLAGS.data_dir = "xyz"
    expected_flags += " --data_dir=xyz"
    self.assertEqual(flags_core.get_nondefault_flags_as_str(), expected_flags)

    flags.FLAGS.hooks = ["aaa", "bbb", "ccc"]
    expected_flags += " --hooks=aaa,bbb,ccc"
    self.assertEqual(flags_core.get_nondefault_flags_as_str(), expected_flags)

    flags.FLAGS.stop_threshold = 3.
    expected_flags += " --stop_threshold=3.0"
    self.assertEqual(flags_core.get_nondefault_flags_as_str(), expected_flags)

    flags.FLAGS.use_synthetic_data = True
    expected_flags += " --use_synthetic_data"
    self.assertEqual(flags_core.get_nondefault_flags_as_str(), expected_flags)

    # Assert that explicit setting a flag to its default value does not cause it
    # to appear in the string
    flags.FLAGS.use_synthetic_data = False
    expected_flags = expected_flags[:-len(" --use_synthetic_data")]
    self.assertEqual(flags_core.get_nondefault_flags_as_str(), expected_flags)
160

Taylor Robie's avatar
Taylor Robie committed
161
162
163

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