keras_utils.py 7.91 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.
# ==============================================================================
"""Helper functions for the Keras implementations of models."""

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

21
22
import multiprocessing
import os
23
24
import time

25
from absl import logging
26
import tensorflow as tf
Toby Boyd's avatar
Toby Boyd committed
27
from tensorflow.core.protobuf import rewriter_config_pb2
Toby Boyd's avatar
Toby Boyd committed
28
from tensorflow.python import tf2
29
from tensorflow.python.eager import profiler
30
31
32
33
34
35
36
37
38


class BatchTimestamp(object):
  """A structure to store batch time stamp."""

  def __init__(self, batch_index, timestamp):
    self.batch_index = batch_index
    self.timestamp = timestamp

39
40
41
42
  def __repr__(self):
    return "'BatchTimestamp<batch_index: {}, timestamp: {}>'".format(
        self.batch_index, self.timestamp)

43
44
45
46
47

class TimeHistory(tf.keras.callbacks.Callback):
  """Callback for Keras models."""

  def __init__(self, batch_size, log_steps):
48
    """Callback for logging performance.
Shining Sun's avatar
Shining Sun committed
49

50
51
    Args:
      batch_size: Total batch size.
52
      log_steps: Interval of steps between logging of batch level stats.
53
54
55
56
    """
    self.batch_size = batch_size
    super(TimeHistory, self).__init__()
    self.log_steps = log_steps
57
    self.global_steps = 0
58

59
    # Logs start of step 1 then end of each step based on log_steps interval.
60
61
    self.timestamp_log = []

62
63
64
    # Records the time each epoch takes to run from start to finish of epoch.
    self.epoch_runtime_log = []

65
66
67
  def on_train_end(self, logs=None):
    self.train_finish_time = time.time()

68
69
70
  def on_epoch_begin(self, epoch, logs=None):
    self.epoch_start = time.time()

71
  def on_batch_begin(self, batch, logs=None):
72
73
74
75
76
    self.global_steps += 1
    if self.global_steps == 1:
      self.start_time = time.time()
      self.timestamp_log.append(BatchTimestamp(self.global_steps,
                                               self.start_time))
77
78

  def on_batch_end(self, batch, logs=None):
79
80
    """Records elapse time of the batch and calculates examples per second."""
    if self.global_steps % self.log_steps == 0:
81
82
83
      timestamp = time.time()
      elapsed_time = timestamp - self.start_time
      examples_per_second = (self.batch_size * self.log_steps) / elapsed_time
84
      self.timestamp_log.append(BatchTimestamp(self.global_steps, timestamp))
85
      logging.info(
86
          "BenchmarkMetric: {'global step':%d, 'time_taken': %f,"
87
88
          "'examples_per_second': %f}",
          self.global_steps, elapsed_time, examples_per_second)
89
      self.start_time = timestamp
90

91
92
93
  def on_epoch_end(self, epoch, logs=None):
    epoch_run_time = time.time() - self.epoch_start
    self.epoch_runtime_log.append(epoch_run_time)
94
95
96
    logging.info(
        "BenchmarkMetric: {'epoch':%d, 'time_taken': %f}",
        epoch, epoch_run_time)
97

98

Zongwei Zhou's avatar
Zongwei Zhou committed
99
100
def get_profiler_callback(model_dir, profile_steps, enable_tensorboard,
                          steps_per_epoch):
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
  """Validate profile_steps flag value and return profiler callback."""
  profile_steps_error_message = (
      'profile_steps must be a comma separated pair of positive integers, '
      'specifying the first and last steps to be profiled.'
  )
  try:
    profile_steps = [int(i) for i in profile_steps.split(',')]
  except ValueError:
    raise ValueError(profile_steps_error_message)
  if len(profile_steps) != 2:
    raise ValueError(profile_steps_error_message)
  start_step, stop_step = profile_steps
  if start_step < 0 or start_step > stop_step:
    raise ValueError(profile_steps_error_message)
  if enable_tensorboard:
116
    logging.warning(
117
118
119
120
        'Both TensorBoard and profiler callbacks are used. Note that the '
        'TensorBoard callback profiles the 2nd step (unless otherwise '
        'specified). Please make sure the steps profiled by the two callbacks '
        'do not overlap.')
Zongwei Zhou's avatar
Zongwei Zhou committed
121
  return ProfilerCallback(model_dir, start_step, stop_step, steps_per_epoch)
122
123
124
125
126


class ProfilerCallback(tf.keras.callbacks.Callback):
  """Save profiles in specified step range to log directory."""

Zongwei Zhou's avatar
Zongwei Zhou committed
127
  def __init__(self, log_dir, start_step, stop_step, steps_per_epoch):
128
129
130
131
    super(ProfilerCallback, self).__init__()
    self.log_dir = log_dir
    self.start_step = start_step
    self.stop_step = stop_step
Zongwei Zhou's avatar
Zongwei Zhou committed
132
133
134
135
136
137
138
139
140
141
142
143
    self.start_epoch = start_step // steps_per_epoch
    self.stop_epoch = stop_step // steps_per_epoch
    self.start_step_in_epoch = start_step % steps_per_epoch
    self.stop_step_in_epoch = stop_step % steps_per_epoch
    self.should_start = False
    self.should_stop = False

  def on_epoch_begin(self, epoch, logs=None):
    if epoch == self.start_epoch:
      self.should_start = True
    if epoch == self.stop_epoch:
      self.should_stop = True
144
145

  def on_batch_begin(self, batch, logs=None):
Zongwei Zhou's avatar
Zongwei Zhou committed
146
147
    if batch == self.start_step_in_epoch and self.should_start:
      self.should_start = False
148
      profiler.start()
149
      logging.info('Profiler started at Step %s', self.start_step)
150
151

  def on_batch_end(self, batch, logs=None):
Zongwei Zhou's avatar
Zongwei Zhou committed
152
153
    if batch == self.stop_step_in_epoch and self.should_stop:
      self.should_stop = False
154
155
      results = profiler.stop()
      profiler.save(self.log_dir, results)
156
      logging.info(
157
158
          'Profiler saved profiles for steps between %s and %s to %s',
          self.start_step, self.stop_step, self.log_dir)
Toby Boyd's avatar
Toby Boyd committed
159
160
161


def set_session_config(enable_eager=False,
162
                       enable_xla=False):
Toby Boyd's avatar
Toby Boyd committed
163
164
  """Sets the session config."""
  if is_v2_0():
165
    set_config_v2(enable_xla=enable_xla)
Toby Boyd's avatar
Toby Boyd committed
166
  else:
167
    config = get_config_proto_v1(enable_xla=enable_xla)
Toby Boyd's avatar
Toby Boyd committed
168
169
170
171
172
173
174
    if enable_eager:
      tf.compat.v1.enable_eager_execution(config=config)
    else:
      sess = tf.Session(config=config)
      tf.keras.backend.set_session(sess)


175
def get_config_proto_v1(enable_xla=False):
Toby Boyd's avatar
Toby Boyd committed
176
177
178
179
180
181
182
183
184
  """Return config proto according to flag settings, or None to use default."""
  config = None
  if enable_xla:
    config = tf.compat.v1.ConfigProto()
    config.graph_options.optimizer_options.global_jit_level = (
        tf.OptimizerOptions.ON_2)
  return config


185
def set_config_v2(enable_xla=False):
Toby Boyd's avatar
Toby Boyd committed
186
187
188
189
  """Config eager context according to flag values using TF 2.0 API."""
  if enable_xla:
    tf.config.optimizer.set_jit(True)

Toby Boyd's avatar
Toby Boyd committed
190

Toby Boyd's avatar
Toby Boyd committed
191
192
def is_v2_0():
  """Returns true if using tf 2.0."""
Toby Boyd's avatar
Toby Boyd committed
193
  return tf2.enabled()
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221


def set_gpu_thread_mode_and_count(gpu_thread_mode,
                                  datasets_num_private_threads,
                                  num_gpus, per_gpu_thread_count):
  """Set GPU thread mode and count, and adjust dataset threads count."""
  cpu_count = multiprocessing.cpu_count()
  logging.info('Logical CPU cores: %s', cpu_count)

  # Allocate private thread pool for each GPU to schedule and launch kernels
  per_gpu_thread_count = per_gpu_thread_count or 2
  os.environ['TF_GPU_THREAD_MODE'] = gpu_thread_mode
  os.environ['TF_GPU_THREAD_COUNT'] = str(per_gpu_thread_count)
  logging.info('TF_GPU_THREAD_COUNT: %s',
               os.environ['TF_GPU_THREAD_COUNT'])
  logging.info('TF_GPU_THREAD_MODE: %s',
               os.environ['TF_GPU_THREAD_MODE'])

  # Limit data preprocessing threadpool to CPU cores minus number of total GPU
  # private threads and memory copy threads.
  total_gpu_thread_count = per_gpu_thread_count * num_gpus
  num_runtime_threads = num_gpus
  if not datasets_num_private_threads:
    datasets_num_private_threads = min(
        cpu_count - total_gpu_thread_count - num_runtime_threads,
        num_gpus * 8)
    logging.info('Set datasets_num_private_threads to %s',
                 datasets_num_private_threads)