run_classifier.py 13.9 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
# 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.
# ==============================================================================
15
"""BERT classification finetuning runner in TF 2.x."""
16
17
18
19
20
21
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import json
import math
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
22
import os
23
24
25
26
27
28

from absl import app
from absl import flags
from absl import logging
import tensorflow as tf

29
from official.modeling import model_training_utils
30
from official.modeling import performance
31
from official.nlp import optimization
32
from official.nlp.bert import bert_models
33
from official.nlp.bert import common_flags
34
from official.nlp.bert import configs as bert_configs
35
36
from official.nlp.bert import input_pipeline
from official.nlp.bert import model_saving_utils
37
from official.utils.misc import distribution_utils
38
from official.utils.misc import keras_utils
39

40

41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
flags.DEFINE_enum(
    'mode', 'train_and_eval', ['train_and_eval', 'export_only'],
    'One of {"train_and_eval", "export_only"}. `train_and_eval`: '
    'trains the model and evaluates in the meantime. '
    '`export_only`: will take the latest checkpoint inside '
    'model_dir and export a `SavedModel`.')
flags.DEFINE_string('train_data_path', None,
                    'Path to training data for BERT classifier.')
flags.DEFINE_string('eval_data_path', None,
                    'Path to evaluation data for BERT classifier.')
# Model training specific flags.
flags.DEFINE_string(
    'input_meta_data_path', None,
    'Path to file that contains meta data about input '
    'to be used for training and evaluation.')
flags.DEFINE_integer('train_batch_size', 32, 'Batch size for training.')
57
flags.DEFINE_integer('eval_batch_size', 32, 'Batch size for evaluation.')
58
59

common_flags.define_common_bert_flags()
60
61
62
63

FLAGS = flags.FLAGS


64
def get_loss_fn(num_classes, loss_factor=1.0):
65
66
67
68
69
70
71
72
73
74
75
  """Gets the classification loss function."""

  def classification_loss_fn(labels, logits):
    """Classification loss."""
    labels = tf.squeeze(labels)
    log_probs = tf.nn.log_softmax(logits, axis=-1)
    one_hot_labels = tf.one_hot(
        tf.cast(labels, dtype=tf.int32), depth=num_classes, dtype=tf.float32)
    per_example_loss = -tf.reduce_sum(
        tf.cast(one_hot_labels, dtype=tf.float32) * log_probs, axis=-1)
    loss = tf.reduce_mean(per_example_loss)
76
    loss *= loss_factor
77
78
79
80
81
    return loss

  return classification_loss_fn


Hongkun Yu's avatar
Hongkun Yu committed
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
def get_dataset_fn(input_file_pattern, max_seq_length, global_batch_size,
                   is_training):
  """Gets a closure to create a dataset."""

  def _dataset_fn(ctx=None):
    """Returns tf.data.Dataset for distributed BERT pretraining."""
    batch_size = ctx.get_per_replica_batch_size(
        global_batch_size) if ctx else global_batch_size
    dataset = input_pipeline.create_classifier_dataset(
        input_file_pattern,
        max_seq_length,
        batch_size,
        is_training=is_training,
        input_pipeline_context=ctx)
    return dataset

  return _dataset_fn


A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
101
102
103
104
105
106
107
108
109
110
111
def run_bert_classifier(strategy,
                        bert_config,
                        input_meta_data,
                        model_dir,
                        epochs,
                        steps_per_epoch,
                        steps_per_loop,
                        eval_steps,
                        warmup_steps,
                        initial_lr,
                        init_checkpoint,
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
112
113
                        train_input_fn,
                        eval_input_fn,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
114
                        custom_callbacks=None,
115
116
                        run_eagerly=False,
                        use_keras_compile_fit=False):
117
118
119
120
121
  """Run BERT classifier training using low-level API."""
  max_seq_length = input_meta_data['max_seq_length']
  num_classes = input_meta_data['num_labels']

  def _get_classifier_model():
122
    """Gets a classifier model."""
123
    classifier_model, core_model = (
124
125
126
127
        bert_models.classifier_model(
            bert_config,
            num_classes,
            max_seq_length,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
128
129
            hub_module_url=FLAGS.hub_module_url,
            hub_module_trainable=FLAGS.hub_module_trainable))
130
    optimizer = optimization.create_optimizer(
131
        initial_lr, steps_per_epoch * epochs, warmup_steps)
132
133
134
135
    classifier_model.optimizer = performance.configure_optimizer(
        optimizer,
        use_float16=common_flags.use_float16(),
        use_graph_rewrite=common_flags.use_graph_rewrite())
136
137
    return classifier_model, core_model

138
139
140
141
142
143
144
145
146
147
148
  # During distributed training, loss used for gradient computation is
  # summed over from all replicas. When Keras compile/fit() API is used,
  # the fit() API internally normalizes the loss by dividing the loss by
  # the number of replicas used for computation. However, when custom
  # training loop is used this is not done automatically and should be
  # done manually by the end user.
  loss_multiplier = 1.0
  if FLAGS.scale_loss and not use_keras_compile_fit:
    loss_multiplier = 1.0 / strategy.num_replicas_in_sync

  loss_fn = get_loss_fn(num_classes, loss_factor=loss_multiplier)
149
150
151
152
153
154
155

  # Defines evaluation metrics function, which will create metrics in the
  # correct device and strategy scope.
  def metric_fn():
    return tf.keras.metrics.SparseCategoricalAccuracy(
        'test_accuracy', dtype=tf.float32)

156
  if use_keras_compile_fit:
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
157
158
    # Start training using Keras compile/fit API.
    logging.info('Training using TF 2.0 Keras compile/fit API with '
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
159
                 'distribution strategy.')
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
160
161
162
163
164
165
166
167
168
169
170
171
    return run_keras_compile_fit(
        model_dir,
        strategy,
        _get_classifier_model,
        train_input_fn,
        eval_input_fn,
        loss_fn,
        metric_fn,
        init_checkpoint,
        epochs,
        steps_per_epoch,
        eval_steps,
172
        custom_callbacks=custom_callbacks)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
173
174
175

  # Use user-defined loop to start training.
  logging.info('Training using customized training loop TF 2.0 with '
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
176
               'distribution strategy.')
177
178
179
180
181
182
  return model_training_utils.run_customized_training_loop(
      strategy=strategy,
      model_fn=_get_classifier_model,
      loss_fn=loss_fn,
      model_dir=model_dir,
      steps_per_epoch=steps_per_epoch,
183
      steps_per_loop=steps_per_loop,
184
185
186
187
188
189
      epochs=epochs,
      train_input_fn=train_input_fn,
      eval_input_fn=eval_input_fn,
      eval_steps=eval_steps,
      init_checkpoint=init_checkpoint,
      metric_fn=metric_fn,
190
191
      custom_callbacks=custom_callbacks,
      run_eagerly=run_eagerly)
192
193


A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
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
def run_keras_compile_fit(model_dir,
                          strategy,
                          model_fn,
                          train_input_fn,
                          eval_input_fn,
                          loss_fn,
                          metric_fn,
                          init_checkpoint,
                          epochs,
                          steps_per_epoch,
                          eval_steps,
                          custom_callbacks=None):
  """Runs BERT classifier model using Keras compile/fit API."""

  with strategy.scope():
    training_dataset = train_input_fn()
    evaluation_dataset = eval_input_fn()
    bert_model, sub_model = model_fn()
    optimizer = bert_model.optimizer

    if init_checkpoint:
      checkpoint = tf.train.Checkpoint(model=sub_model)
      checkpoint.restore(init_checkpoint).assert_existing_objects_matched()

    bert_model.compile(optimizer=optimizer, loss=loss_fn, metrics=[metric_fn()])

220
221
222
223
224
    summary_dir = os.path.join(model_dir, 'summaries')
    summary_callback = tf.keras.callbacks.TensorBoard(summary_dir)
    checkpoint_path = os.path.join(model_dir, 'checkpoint')
    checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
        checkpoint_path, save_weights_only=True)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241

    if custom_callbacks is not None:
      custom_callbacks += [summary_callback, checkpoint_callback]
    else:
      custom_callbacks = [summary_callback, checkpoint_callback]

    bert_model.fit(
        x=training_dataset,
        validation_data=evaluation_dataset,
        steps_per_epoch=steps_per_epoch,
        epochs=epochs,
        validation_steps=eval_steps,
        callbacks=custom_callbacks)

    return bert_model


242
def export_classifier(model_export_path, input_meta_data,
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
243
244
                      restore_model_using_load_weights,
                      bert_config, model_dir):
245
246
247
248
249
  """Exports a trained model as a `SavedModel` for inference.

  Args:
    model_export_path: a string specifying the path to the SavedModel directory.
    input_meta_data: dictionary containing meta data about input and model.
250
251
252
253
254
255
    restore_model_using_load_weights: Whether to use checkpoint.restore() API
      for custom checkpoint or to use model.load_weights() API.
      There are 2 different ways to save checkpoints. One is using
      tf.train.Checkpoint and another is using Keras model.save_weights().
      Custom training loop implementation uses tf.train.Checkpoint API
      and Keras ModelCheckpoint callback internally uses model.save_weights()
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
256
      API. Since these two API's cannot be used together, model loading logic
257
      must be take into account how model checkpoint was saved.
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
258
259
260
    bert_config: Bert configuration file to define core bert layers.
    model_dir: The directory where the model weights and training/evaluation
      summaries are stored.
261
262
263
264
265
266

  Raises:
    Export path is not specified, got an empty string or None.
  """
  if not model_export_path:
    raise ValueError('Export path is not specified: %s' % model_export_path)
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
267
268
  if not model_dir:
    raise ValueError('Export path is not specified: %s' % model_dir)
269

Zongwei Zhou's avatar
Zongwei Zhou committed
270
271
  # Export uses float32 for now, even if training uses mixed precision.
  tf.keras.mixed_precision.experimental.set_policy('float32')
272
  classifier_model = bert_models.classifier_model(
Zongwei Zhou's avatar
Zongwei Zhou committed
273
      bert_config, input_meta_data['num_labels'],
274
      input_meta_data['max_seq_length'])[0]
275

276
  model_saving_utils.export_bert_model(
277
278
      model_export_path,
      model=classifier_model,
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
279
      checkpoint_dir=model_dir,
280
      restore_model_using_load_weights=restore_model_using_load_weights)
281
282


Hongkun Yu's avatar
Hongkun Yu committed
283
284
def run_bert(strategy,
             input_meta_data,
285
             model_config,
Hongkun Yu's avatar
Hongkun Yu committed
286
287
             train_input_fn=None,
             eval_input_fn=None):
288
289
  """Run BERT training."""
  if FLAGS.mode == 'export_only':
290
291
292
293
    # As Keras ModelCheckpoint callback used with Keras compile/fit() API
    # internally uses model.save_weights() to save checkpoints, we must
    # use model.load_weights() when Keras compile/fit() is used.
    export_classifier(FLAGS.model_export_path, input_meta_data,
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
294
                      FLAGS.use_keras_compile_fit,
295
                      model_config, FLAGS.model_dir)
296
297
298
299
    return

  if FLAGS.mode != 'train_and_eval':
    raise ValueError('Unsupported mode is specified: %s' % FLAGS.mode)
300
301
  # Enables XLA in Session Config. Should not be set for TPU.
  keras_utils.set_config_v2(FLAGS.enable_xla)
302
  performance.set_mixed_precision_policy(common_flags.dtype())
303
304
305
306
307
308
309
310
311
312

  epochs = FLAGS.num_train_epochs
  train_data_size = input_meta_data['train_data_size']
  steps_per_epoch = int(train_data_size / FLAGS.train_batch_size)
  warmup_steps = int(epochs * train_data_size * 0.1 / FLAGS.train_batch_size)
  eval_steps = int(
      math.ceil(input_meta_data['eval_data_size'] / FLAGS.eval_batch_size))

  if not strategy:
    raise ValueError('Distribution strategy has not been specified.')
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
313

314
315
316
317
318
319
320
321
322
  if FLAGS.log_steps:
    custom_callbacks = [keras_utils.TimeHistory(
        batch_size=FLAGS.train_batch_size,
        log_steps=FLAGS.log_steps,
        logdir=FLAGS.model_dir,
    )]
  else:
    custom_callbacks = None

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
323
  trained_model = run_bert_classifier(
324
      strategy,
325
      model_config,
326
327
328
329
      input_meta_data,
      FLAGS.model_dir,
      epochs,
      steps_per_epoch,
330
      FLAGS.steps_per_loop,
331
332
333
334
      eval_steps,
      warmup_steps,
      FLAGS.learning_rate,
      FLAGS.init_checkpoint,
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
335
336
      train_input_fn,
      eval_input_fn,
337
      run_eagerly=FLAGS.run_eagerly,
338
339
      use_keras_compile_fit=FLAGS.use_keras_compile_fit,
      custom_callbacks=custom_callbacks)
340

341
  if FLAGS.model_export_path:
342
343
344
    # As Keras ModelCheckpoint callback used with Keras compile/fit() API
    # internally uses model.save_weights() to save checkpoints, we must
    # use model.load_weights() when Keras compile/fit() is used.
345
    model_saving_utils.export_bert_model(
346
347
348
        FLAGS.model_export_path,
        model=trained_model,
        restore_model_using_load_weights=FLAGS.use_keras_compile_fit)
349
350
  return trained_model

351
352
353
354

def main(_):
  # Users should always run this script under TF 2.x
  assert tf.version.VERSION.startswith('2.')
355

356
357
358
359
360
361
  with tf.io.gfile.GFile(FLAGS.input_meta_data_path, 'rb') as reader:
    input_meta_data = json.loads(reader.read().decode('utf-8'))

  if not FLAGS.model_dir:
    FLAGS.model_dir = '/tmp/bert20/'

362
363
364
365
  strategy = distribution_utils.get_distribution_strategy(
      distribution_strategy=FLAGS.distribution_strategy,
      num_gpus=FLAGS.num_gpus,
      tpu_address=FLAGS.tpu)
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
366
  max_seq_length = input_meta_data['max_seq_length']
Hongkun Yu's avatar
Hongkun Yu committed
367
  train_input_fn = get_dataset_fn(
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
368
      FLAGS.train_data_path,
Hongkun Yu's avatar
Hongkun Yu committed
369
370
371
372
      max_seq_length,
      FLAGS.train_batch_size,
      is_training=True)
  eval_input_fn = get_dataset_fn(
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
373
      FLAGS.eval_data_path,
Hongkun Yu's avatar
Hongkun Yu committed
374
375
376
377
      max_seq_length,
      FLAGS.eval_batch_size,
      is_training=False)

378
379
380
  bert_config = bert_configs.BertConfig.from_json_file(FLAGS.bert_config_file)
  run_bert(strategy, input_meta_data, bert_config, train_input_fn,
           eval_input_fn)
381
382
383
384
385


if __name__ == '__main__':
  flags.mark_flag_as_required('bert_config_file')
  flags.mark_flag_as_required('input_meta_data_path')
386
  flags.mark_flag_as_required('model_dir')
387
  app.run(main)