misc.py 10.5 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# Copyright 2019 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.
# ==============================================================================
"""Misc for Transformer."""

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

Toby Boyd's avatar
Toby Boyd committed
21
# pylint: disable=g-bad-import-order
22
from absl import flags
Toby Boyd's avatar
Toby Boyd committed
23
import tensorflow as tf
24

25
26
27
28
# TODO(tianlin) Import internal library. Remove this when some functions for
# different TF versions are fixed.
from tensorflow.python import tf2 as tf2_internal

29
from official.nlp.transformer import model_params
30
from official.utils.flags import core as flags_core
Toby Boyd's avatar
Toby Boyd committed
31
32
33
from official.utils.misc import keras_utils

FLAGS = flags.FLAGS
34
35

PARAMS_MAP = {
Toby Boyd's avatar
Toby Boyd committed
36
37
38
    'tiny': model_params.TINY_PARAMS,
    'base': model_params.BASE_PARAMS,
    'big': model_params.BIG_PARAMS,
39
40
41
}


42
43
44
45
46
def is_v2():
  """Returns whether it is v2."""
  return tf2_internal.enabled()


47
48
49
def get_model_params(param_set, num_gpus):
  """Gets predefined model params."""
  if num_gpus > 1:
Toby Boyd's avatar
Toby Boyd committed
50
    if param_set == 'big':
51
      return model_params.BIG_MULTI_GPU_PARAMS.copy()
Toby Boyd's avatar
Toby Boyd committed
52
    elif param_set == 'base':
53
54
      return model_params.BASE_MULTI_GPU_PARAMS.copy()
    else:
Toby Boyd's avatar
Toby Boyd committed
55
      raise ValueError('Not valid params: param_set={} num_gpus={}'.format(
56
57
58
59
60
61
62
          param_set, num_gpus))

  return PARAMS_MAP[param_set].copy()


def define_transformer_flags():
  """Add flags and flag validators for running transformer_main."""
63
  # Add common flags (data_dir, model_dir, etc.).
64
  flags_core.define_base(num_gpu=True, distribution_strategy=True)
65
66
67
68
69
70
  flags_core.define_performance(
      num_parallel_calls=True,
      inter_op=False,
      intra_op=False,
      synthetic_data=True,
      max_train_steps=False,
71
72
      dtype=True,
      loss_scale=True,
Toby Boyd's avatar
Toby Boyd committed
73
      all_reduce_alg=True,
74
75
76
      num_packs=True,
      tf_gpu_thread_mode=True,
      datasets_num_private_threads=True,
77
      enable_xla=True,
Vinh Nguyen's avatar
Vinh Nguyen committed
78
      fp16_implementation=True
79
  )
Toby Boyd's avatar
Toby Boyd committed
80
81
82
83
84
85
86
87
88
89
90

  # Additional performance flags
  # TODO(b/76028325): Remove when generic layout optimizer is ready.
  flags.DEFINE_boolean(
      name='enable_grappler_layout_optimizer',
      default=True,
      help='Enable Grappler layout optimizer. Currently Grappler can '
           'de-optimize fp16 graphs by forcing NCHW layout for all '
           'convolutions and batch normalizations, and this flag allows to '
           'disable it.'
  )
91

92
93
94
  flags_core.define_benchmark()
  flags_core.define_device(tpu=True)

Toby Boyd's avatar
Toby Boyd committed
95
  flags.DEFINE_integer(
Toby Boyd's avatar
Toby Boyd committed
96
      name='train_steps', short_name='ts', default=300000,
Toby Boyd's avatar
Toby Boyd committed
97
98
      help=flags_core.help_wrap('The number of steps used to train.'))
  flags.DEFINE_integer(
99
      name='steps_between_evals', short_name='sbe', default=5000,
Toby Boyd's avatar
Toby Boyd committed
100
101
102
      help=flags_core.help_wrap(
          'The Number of training steps to run between evaluations. This is '
          'used if --train_steps is defined.'))
103
104
105
  flags.DEFINE_boolean(
      name='enable_time_history', default=True,
      help='Whether to enable TimeHistory callback.')
Toby Boyd's avatar
Toby Boyd committed
106
107
108
  flags.DEFINE_boolean(
      name='enable_tensorboard', default=False,
      help='Whether to enable Tensorboard callback.')
109
110
111
  flags.DEFINE_boolean(
      name='enable_metrics_in_training', default=False,
      help='Whether to enable metrics during training.')
112
113
114
115
  flags.DEFINE_boolean(
      name='enable_mlir_bridge',
      default=False,
      help='Whether to enable the TF to XLA bridge.')
Toby Boyd's avatar
Toby Boyd committed
116
  # Set flags from the flags_core module as 'key flags' so they're listed when
117
118
119
120
121
122
  # the '-h' flag is used. Without this line, the flags defined above are
  # only shown in the full `--helpful` help text.
  flags.adopt_module_key_flags(flags_core)

  # Add transformer-specific flags
  flags.DEFINE_enum(
Toby Boyd's avatar
Toby Boyd committed
123
      name='param_set', short_name='mp', default='big',
124
125
      enum_values=PARAMS_MAP.keys(),
      help=flags_core.help_wrap(
Toby Boyd's avatar
Toby Boyd committed
126
127
128
129
130
131
          'Parameter set to use when creating and training the model. The '
          'parameters define the input shape (batch size and max length), '
          'model configuration (size of embedding, # of hidden layers, etc.), '
          'and various other settings. The big parameter set increases the '
          'default batch size, embedding/hidden size, and filter size. For a '
          'complete list of parameters, please see model/model_params.py.'))
132
133

  flags.DEFINE_bool(
134
      name='static_batch', short_name='sb', default=False,
135
      help=flags_core.help_wrap(
Toby Boyd's avatar
Toby Boyd committed
136
137
138
139
140
141
          'Whether the batches in the dataset should have static shapes. In '
          'general, this setting should be False. Dynamic shapes allow the '
          'inputs to be grouped so that the number of padding tokens is '
          'minimized, and helps model training. In cases where the input shape '
          'must be static (e.g. running on TPU), this setting will be ignored '
          'and static batching will always be used.'))
142
143
144
145
146
147
  flags.DEFINE_integer(
      name='max_length', short_name='ml', default=256,
      help=flags_core.help_wrap(
          'Max sentence length for Transformer. Default is 256. Note: Usually '
          'it is more effective to use a smaller max length if static_batch is '
          'enabled, e.g. 64.'))
148
149
150

  # Flags for training with steps (may be used for debugging)
  flags.DEFINE_integer(
Toby Boyd's avatar
Toby Boyd committed
151
152
      name='validation_steps', short_name='vs', default=64,
      help=flags_core.help_wrap('The number of steps used in validation.'))
153
154
155

  # BLEU score computation
  flags.DEFINE_string(
Toby Boyd's avatar
Toby Boyd committed
156
      name='bleu_source', short_name='bls', default=None,
157
      help=flags_core.help_wrap(
Toby Boyd's avatar
Toby Boyd committed
158
159
          'Path to source file containing text translate when calculating the '
          'official BLEU score. Both --bleu_source and --bleu_ref must be set. '
160
          ))
161
  flags.DEFINE_string(
Toby Boyd's avatar
Toby Boyd committed
162
      name='bleu_ref', short_name='blr', default=None,
163
      help=flags_core.help_wrap(
Toby Boyd's avatar
Toby Boyd committed
164
165
          'Path to source file containing text translate when calculating the '
          'official BLEU score. Both --bleu_source and --bleu_ref must be set. '
166
          ))
167
  flags.DEFINE_string(
Toby Boyd's avatar
Toby Boyd committed
168
      name='vocab_file', short_name='vf', default=None,
169
      help=flags_core.help_wrap(
Toby Boyd's avatar
Toby Boyd committed
170
171
172
          'Path to subtoken vocabulary file. If data_download.py was used to '
          'download and encode the training data, look in the data_dir to find '
          'the vocab file.'))
173
  flags.DEFINE_string(
Toby Boyd's avatar
Toby Boyd committed
174
175
      name='mode', default='train',
      help=flags_core.help_wrap('mode: train, eval, or predict'))
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
176
177
178
179
180
  flags.DEFINE_bool(
      name='use_ctl',
      default=False,
      help=flags_core.help_wrap(
          'Whether the model runs with custom training loop.'))
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
  flags.DEFINE_integer(
      name='decode_batch_size',
      default=32,
      help=flags_core.help_wrap(
          'Global batch size used for Transformer autoregressive decoding on '
          'TPU.'))
  flags.DEFINE_integer(
      name='decode_max_length',
      default=97,
      help=flags_core.help_wrap(
          'Max sequence length of the decode/eval data. This is used by '
          'Transformer autoregressive decoding on TPU to have minimum '
          'paddings.'))
  flags.DEFINE_bool(
      name='padded_decode',
      default=False,
      help=flags_core.help_wrap(
          'Whether the autoregressive decoding runs with input data padded to '
          'the decode_max_length. For TPU/XLA-GPU runs, this flag has to be '
          'set due the static shape requirement. Although CPU/GPU could also '
          'use padded_decode, it has not been tested. In addition, this method '
          'will introduce unnecessary overheads which grow quadratically with '
          'the max sequence length.'))
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
204
205
206
207
208
209
  flags.DEFINE_bool(
      name='enable_checkpointing',
      default=True,
      help=flags_core.help_wrap(
          'Whether to do checkpointing during training. When running under '
          'benchmark harness, we will avoid checkpointing.'))
210

Toby Boyd's avatar
Toby Boyd committed
211
212
  flags_core.set_defaults(data_dir='/tmp/translate_ende',
                          model_dir='/tmp/transformer_model',
213
                          batch_size=None)
214
215
216

  # pylint: disable=unused-variable
  @flags.multi_flags_validator(
Toby Boyd's avatar
Toby Boyd committed
217
218
      ['bleu_source', 'bleu_ref'],
      message='Both or neither --bleu_source and --bleu_ref must be defined.')
219
  def _check_bleu_files(flags_dict):
Toby Boyd's avatar
Toby Boyd committed
220
221
    return (flags_dict['bleu_source'] is None) == (
        flags_dict['bleu_ref'] is None)
222
223

  @flags.multi_flags_validator(
Toby Boyd's avatar
Toby Boyd committed
224
225
226
      ['bleu_source', 'bleu_ref', 'vocab_file'],
      message='--vocab_file must be defined if --bleu_source and --bleu_ref '
              'are defined.')
227
  def _check_bleu_vocab_file(flags_dict):
Toby Boyd's avatar
Toby Boyd committed
228
229
    if flags_dict['bleu_source'] and flags_dict['bleu_ref']:
      return flags_dict['vocab_file'] is not None
230
231
232
    return True
  # pylint: enable=unused-variable

Toby Boyd's avatar
Toby Boyd committed
233

234
def get_callbacks():
Toby Boyd's avatar
Toby Boyd committed
235
236
  """Returns common callbacks."""
  callbacks = []
237
  if FLAGS.enable_time_history:
Will Cromar's avatar
Will Cromar committed
238
239
240
241
    time_callback = keras_utils.TimeHistory(
        FLAGS.batch_size,
        FLAGS.log_steps,
        FLAGS.model_dir if FLAGS.enable_tensorboard else None)
242
    callbacks.append(time_callback)
Toby Boyd's avatar
Toby Boyd committed
243
244
245
246
247
248
249
250
251

  if FLAGS.enable_tensorboard:
    tensorboard_callback = tf.keras.callbacks.TensorBoard(
        log_dir=FLAGS.model_dir)
    callbacks.append(tensorboard_callback)

  return callbacks


Tayo Oguntebi's avatar
Tayo Oguntebi committed
252
253
def update_stats(history, stats, callbacks):
  """Normalizes and updates dictionary of stats.
Toby Boyd's avatar
Toby Boyd committed
254
255
256

  Args:
    history: Results of the training step.
Tayo Oguntebi's avatar
Tayo Oguntebi committed
257
    stats: Dict with pre-existing training stats.
Toby Boyd's avatar
Toby Boyd committed
258
259
260
261
262
263
264
    callbacks: a list of callbacks which might include a time history callback
      used during keras.fit.
  """

  if history and history.history:
    train_hist = history.history
    # Gets final loss from training.
265
    stats['loss'] = float(train_hist['loss'][-1])
Toby Boyd's avatar
Toby Boyd committed
266
267

  if not callbacks:
Tayo Oguntebi's avatar
Tayo Oguntebi committed
268
    return
Toby Boyd's avatar
Toby Boyd committed
269
270
271
272
273
274
275
276
277
278
279
280

  # Look for the time history callback which was used during keras.fit
  for callback in callbacks:
    if isinstance(callback, keras_utils.TimeHistory):
      timestamp_log = callback.timestamp_log
      stats['step_timestamp_log'] = timestamp_log
      stats['train_finish_time'] = callback.train_finish_time
      if len(timestamp_log) > 1:
        stats['avg_exp_per_second'] = (
            callback.batch_size * callback.log_steps *
            (len(callback.timestamp_log)-1) /
            (timestamp_log[-1].timestamp - timestamp_log[0].timestamp))