logger_test.py 11.3 KB
Newer Older
Scott Zhu's avatar
Scott Zhu committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
# 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.
# ==============================================================================

"""Tests for benchmark logger."""

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

import json
import os
import tempfile
25
import time
26
import unittest
Scott Zhu's avatar
Scott Zhu committed
27

28
29
import mock
from absl.testing import flagsaver
Karmel Allison's avatar
Karmel Allison committed
30
import tensorflow as tf  # pylint: disable=g-bad-import-order
31
from absl import logging
Scott Zhu's avatar
Scott Zhu committed
32

33
34
35
36
37
try:
  from google.cloud import bigquery
except ImportError:
  bigquery = None

38
from official.utils.misc import keras_utils
39
from official.utils.flags import core as flags_core
40
from official.utils.logs import logger
Scott Zhu's avatar
Scott Zhu committed
41
42
43
44


class BenchmarkLoggerTest(tf.test.TestCase):

45
46
47
48
49
  @classmethod
  def setUpClass(cls):  # pylint: disable=invalid-name
    super(BenchmarkLoggerTest, cls).setUpClass()
    flags_core.define_benchmark()

Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
50
  def test_get_default_benchmark_logger(self):
51
    with flagsaver.flagsaver(benchmark_logger_type="foo"):
52
53
      self.assertIsInstance(logger.get_benchmark_logger(),
                            logger.BaseBenchmarkLogger)
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
54
55

  def test_config_base_benchmark_logger(self):
56
    with flagsaver.flagsaver(benchmark_logger_type="BaseBenchmarkLogger"):
57
58
59
      logger.config_benchmark_logger()
      self.assertIsInstance(logger.get_benchmark_logger(),
                            logger.BaseBenchmarkLogger)
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
60
61

  def test_config_benchmark_file_logger(self):
62
63
    # Set the benchmark_log_dir first since the benchmark_logger_type will need
    # the value to be set when it does the validation.
64
65
    with flagsaver.flagsaver(benchmark_log_dir="/tmp"):
      with flagsaver.flagsaver(benchmark_logger_type="BenchmarkFileLogger"):
66
67
68
69
        logger.config_benchmark_logger()
        self.assertIsInstance(logger.get_benchmark_logger(),
                              logger.BenchmarkFileLogger)

70
71
72
73
74
  @mock.patch("official.utils.logs.logger.config_benchmark_logger")
  def test_benchmark_context(self, mock_config_benchmark_logger):
    mock_logger = mock.MagicMock()
    mock_config_benchmark_logger.return_value = mock_logger
    with logger.benchmark_context(None):
75
      logging.info("start benchmarking")
76
77
78
79
80
81
82
83
84
85
86
    mock_logger.on_finish.assert_called_once_with(logger.RUN_STATUS_SUCCESS)

  @mock.patch("official.utils.logs.logger.config_benchmark_logger")
  def test_benchmark_context_failure(self, mock_config_benchmark_logger):
    mock_logger = mock.MagicMock()
    mock_config_benchmark_logger.return_value = mock_logger
    with self.assertRaises(RuntimeError):
      with logger.benchmark_context(None):
        raise RuntimeError("training error")
    mock_logger.on_finish.assert_called_once_with(logger.RUN_STATUS_FAILURE)

Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
87
88
89
90
91

class BaseBenchmarkLoggerTest(tf.test.TestCase):

  def setUp(self):
    super(BaseBenchmarkLoggerTest, self).setUp()
92
    self._actual_log = logging.info
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
93
94
95
96
97
98
    self.logged_message = None

    def mock_log(*args, **kwargs):
      self.logged_message = args
      self._actual_log(*args, **kwargs)

99
    logging.info = mock_log
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
100
101
102

  def tearDown(self):
    super(BaseBenchmarkLoggerTest, self).tearDown()
103
    logging.info = self._actual_log
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
104
105
106
107
108
109
110
111
112
113
114

  def test_log_metric(self):
    log = logger.BaseBenchmarkLogger()
    log.log_metric("accuracy", 0.999, global_step=1e4, extras={"name": "value"})

    expected_log_prefix = "Benchmark metric:"
    self.assertRegexpMatches(str(self.logged_message), expected_log_prefix)


class BenchmarkFileLoggerTest(tf.test.TestCase):

115
  def setUp(self):
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
116
    super(BenchmarkFileLoggerTest, self).setUp()
117
118
119
120
121
122
    # Avoid pulling extra env vars from test environment which affects the test
    # result, eg. Kokoro test has a TF_PKG env which affect the test case
    # test_collect_tensorflow_environment_variables()
    self.original_environ = dict(os.environ)
    os.environ.clear()

Scott Zhu's avatar
Scott Zhu committed
123
  def tearDown(self):
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
124
    super(BenchmarkFileLoggerTest, self).tearDown()
125
    tf.io.gfile.rmtree(self.get_temp_dir())
126
127
    os.environ.clear()
    os.environ.update(self.original_environ)
Scott Zhu's avatar
Scott Zhu committed
128
129
130

  def test_create_logging_dir(self):
    non_exist_temp_dir = os.path.join(self.get_temp_dir(), "unknown_dir")
131
    self.assertFalse(tf.io.gfile.isdir(non_exist_temp_dir))
Scott Zhu's avatar
Scott Zhu committed
132

Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
133
    logger.BenchmarkFileLogger(non_exist_temp_dir)
134
    self.assertTrue(tf.io.gfile.isdir(non_exist_temp_dir))
Scott Zhu's avatar
Scott Zhu committed
135
136
137

  def test_log_metric(self):
    log_dir = tempfile.mkdtemp(dir=self.get_temp_dir())
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
138
    log = logger.BenchmarkFileLogger(log_dir)
Scott Zhu's avatar
Scott Zhu committed
139
140
141
    log.log_metric("accuracy", 0.999, global_step=1e4, extras={"name": "value"})

    metric_log = os.path.join(log_dir, "metric.log")
142
143
    self.assertTrue(tf.io.gfile.exists(metric_log))
    with tf.io.gfile.GFile(metric_log) as f:
Scott Zhu's avatar
Scott Zhu committed
144
145
146
147
148
      metric = json.loads(f.readline())
      self.assertEqual(metric["name"], "accuracy")
      self.assertEqual(metric["value"], 0.999)
      self.assertEqual(metric["unit"], None)
      self.assertEqual(metric["global_step"], 1e4)
149
      self.assertEqual(metric["extras"], [{"name": "name", "value": "value"}])
Scott Zhu's avatar
Scott Zhu committed
150
151
152

  def test_log_multiple_metrics(self):
    log_dir = tempfile.mkdtemp(dir=self.get_temp_dir())
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
153
    log = logger.BenchmarkFileLogger(log_dir)
Scott Zhu's avatar
Scott Zhu committed
154
155
156
157
    log.log_metric("accuracy", 0.999, global_step=1e4, extras={"name": "value"})
    log.log_metric("loss", 0.02, global_step=1e4)

    metric_log = os.path.join(log_dir, "metric.log")
158
159
    self.assertTrue(tf.io.gfile.exists(metric_log))
    with tf.io.gfile.GFile(metric_log) as f:
Scott Zhu's avatar
Scott Zhu committed
160
161
162
163
164
      accuracy = json.loads(f.readline())
      self.assertEqual(accuracy["name"], "accuracy")
      self.assertEqual(accuracy["value"], 0.999)
      self.assertEqual(accuracy["unit"], None)
      self.assertEqual(accuracy["global_step"], 1e4)
165
      self.assertEqual(accuracy["extras"], [{"name": "name", "value": "value"}])
Scott Zhu's avatar
Scott Zhu committed
166
167
168
169
170
171

      loss = json.loads(f.readline())
      self.assertEqual(loss["name"], "loss")
      self.assertEqual(loss["value"], 0.02)
      self.assertEqual(loss["unit"], None)
      self.assertEqual(loss["global_step"], 1e4)
172
      self.assertEqual(loss["extras"], [])
Scott Zhu's avatar
Scott Zhu committed
173

Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
174
  def test_log_non_number_value(self):
Scott Zhu's avatar
Scott Zhu committed
175
    log_dir = tempfile.mkdtemp(dir=self.get_temp_dir())
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
176
    log = logger.BenchmarkFileLogger(log_dir)
Scott Zhu's avatar
Scott Zhu committed
177
178
179
180
    const = tf.constant(1)
    log.log_metric("accuracy", const)

    metric_log = os.path.join(log_dir, "metric.log")
181
    self.assertFalse(tf.io.gfile.exists(metric_log))
Scott Zhu's avatar
Scott Zhu committed
182

183
  def test_log_evaluation_result(self):
184
185
186
    eval_result = {"loss": 0.46237424,
                   "global_step": 207082,
                   "accuracy": 0.9285}
187
    log_dir = tempfile.mkdtemp(dir=self.get_temp_dir())
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
188
189
    log = logger.BenchmarkFileLogger(log_dir)
    log.log_evaluation_result(eval_result)
190
191

    metric_log = os.path.join(log_dir, "metric.log")
192
193
    self.assertTrue(tf.io.gfile.exists(metric_log))
    with tf.io.gfile.GFile(metric_log) as f:
194
195
196
197
198
199
      accuracy = json.loads(f.readline())
      self.assertEqual(accuracy["name"], "accuracy")
      self.assertEqual(accuracy["value"], 0.9285)
      self.assertEqual(accuracy["unit"], None)
      self.assertEqual(accuracy["global_step"], 207082)

200
201
202
203
204
205
      loss = json.loads(f.readline())
      self.assertEqual(loss["name"], "loss")
      self.assertEqual(loss["value"], 0.46237424)
      self.assertEqual(loss["unit"], None)
      self.assertEqual(loss["global_step"], 207082)

206
207
208
  def test_log_evaluation_result_with_invalid_type(self):
    eval_result = "{'loss': 0.46237424, 'global_step': 207082}"
    log_dir = tempfile.mkdtemp(dir=self.get_temp_dir())
Qianli Scott Zhu's avatar
Qianli Scott Zhu committed
209
210
    log = logger.BenchmarkFileLogger(log_dir)
    log.log_evaluation_result(eval_result)
211
212

    metric_log = os.path.join(log_dir, "metric.log")
213
    self.assertFalse(tf.io.gfile.exists(metric_log))
214

215
216
217
218
219
220
221
222
223
224
225
  @mock.patch("official.utils.logs.logger._gather_run_info")
  def test_log_run_info(self, mock_gather_run_info):
    log_dir = tempfile.mkdtemp(dir=self.get_temp_dir())
    log = logger.BenchmarkFileLogger(log_dir)
    run_info = {"model_name": "model_name",
                "dataset": "dataset_name",
                "run_info": "run_value"}
    mock_gather_run_info.return_value = run_info
    log.log_run_info("model_name", "dataset_name", {})

    run_log = os.path.join(log_dir, "benchmark_run.log")
226
227
    self.assertTrue(tf.io.gfile.exists(run_log))
    with tf.io.gfile.GFile(run_log) as f:
228
229
230
231
232
      run_info = json.loads(f.readline())
      self.assertEqual(run_info["model_name"], "model_name")
      self.assertEqual(run_info["dataset"], "dataset_name")
      self.assertEqual(run_info["run_info"], "run_value")

233
234
235
236
  def test_collect_tensorflow_info(self):
    run_info = {}
    logger._collect_tensorflow_info(run_info)
    self.assertNotEqual(run_info["tensorflow_version"], {})
237
238
239
240
    self.assertEqual(run_info["tensorflow_version"]["version"],
                     tf.version.VERSION)
    self.assertEqual(run_info["tensorflow_version"]["git_hash"],
                     tf.version.GIT_VERSION)
241

242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
  def test_collect_run_params(self):
    run_info = {}
    run_parameters = {
        "batch_size": 32,
        "synthetic_data": True,
        "train_epochs": 100.00,
        "dtype": "fp16",
        "resnet_size": 50,
        "random_tensor": tf.constant(2.0)
    }
    logger._collect_run_params(run_info, run_parameters)
    self.assertEqual(len(run_info["run_parameters"]), 6)
    self.assertEqual(run_info["run_parameters"][0],
                     {"name": "batch_size", "long_value": 32})
    self.assertEqual(run_info["run_parameters"][1],
                     {"name": "dtype", "string_value": "fp16"})
258
259
260
261
262
263
    v1_tensor = {"name": "random_tensor", "string_value":
                     "Tensor(\"Const:0\", shape=(), dtype=float32)"}
    v2_tensor = {"name": "random_tensor", "string_value":
                     "tf.Tensor(2.0, shape=(), dtype=float32)"}
    self.assertIn(run_info["run_parameters"][2], [v1_tensor, v2_tensor])

264

265
266
267
268
269
270
271
    self.assertEqual(run_info["run_parameters"][3],
                     {"name": "resnet_size", "long_value": 50})
    self.assertEqual(run_info["run_parameters"][4],
                     {"name": "synthetic_data", "bool_value": "True"})
    self.assertEqual(run_info["run_parameters"][5],
                     {"name": "train_epochs", "float_value": 100.00})

272
273
  def test_collect_tensorflow_environment_variables(self):
    os.environ["TF_ENABLE_WINOGRAD_NONFUSED"] = "1"
274
275
    os.environ["TF_OTHER"] = "2"
    os.environ["OTHER"] = "3"
276
277
278
279

    run_info = {}
    logger._collect_tensorflow_environment_variables(run_info)
    self.assertIsNotNone(run_info["tensorflow_environment_variables"])
280
281
282
283
284
285
    expected_tf_envs = [
        {"name": "TF_ENABLE_WINOGRAD_NONFUSED", "value": "1"},
        {"name": "TF_OTHER", "value": "2"},
    ]
    self.assertEqual(run_info["tensorflow_environment_variables"],
                     expected_tf_envs)
286
287

  def test_collect_memory_info(self):
288
    run_info = {"machine_config": {}}
289
    logger._collect_memory_info(run_info)
290
291
    self.assertIsNotNone(run_info["machine_config"]["memory_total"])
    self.assertIsNotNone(run_info["machine_config"]["memory_available"])
292

293

Scott Zhu's avatar
Scott Zhu committed
294
295
if __name__ == "__main__":
  tf.test.main()