run_classifier.py 16.4 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

from absl import app
from absl import flags
from absl import logging
Le Hou's avatar
Le Hou committed
27
import gin
28
import tensorflow as tf
29
from official.modeling import performance
30
from official.nlp import optimization
31
from official.nlp.bert import bert_models
32
from official.nlp.bert import common_flags
33
from official.nlp.bert import configs as bert_configs
34
35
from official.nlp.bert import input_pipeline
from official.nlp.bert import model_saving_utils
36
from official.nlp.bert import model_training_utils
37
from official.utils.misc import distribution_utils
38
from official.utils.misc import keras_utils
39
40

flags.DEFINE_enum(
Hongkun Yu's avatar
Hongkun Yu committed
41
42
    'mode', 'train_and_eval', ['train_and_eval', 'export_only', 'predict'],
    'One of {"train_and_eval", "export_only", "predict"}. `train_and_eval`: '
43
44
    'trains the model and evaluates in the meantime. '
    '`export_only`: will take the latest checkpoint inside '
Hongkun Yu's avatar
Hongkun Yu committed
45
46
    'model_dir and export a `SavedModel`. `predict`: takes a checkpoint and '
    'restores the model to output predictions on the test set.')
47
48
49
50
51
52
53
54
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.')
flags.DEFINE_string(
    'input_meta_data_path', None,
    'Path to file that contains meta data about input '
    'to be used for training and evaluation.')
Hongkun Yu's avatar
Hongkun Yu committed
55
56
flags.DEFINE_string('predict_checkpoint_path', None,
                    'Path to the checkpoint for predictions.')
57
flags.DEFINE_integer('train_batch_size', 32, 'Batch size for training.')
58
flags.DEFINE_integer('eval_batch_size', 32, 'Batch size for evaluation.')
59
60

common_flags.define_common_bert_flags()
61
62
63
64

FLAGS = flags.FLAGS


65
def get_loss_fn(num_classes):
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)
76
    return tf.reduce_mean(per_example_loss)
77
78
79
80

  return classification_loss_fn


Hongkun Yu's avatar
Hongkun Yu committed
81
82
83
84
85
86
87
88
89
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(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
90
        tf.io.gfile.glob(input_file_pattern),
Hongkun Yu's avatar
Hongkun Yu committed
91
92
93
94
95
96
97
98
99
        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
100
101
102
103
104
105
106
107
108
109
110
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
111
112
                        train_input_fn,
                        eval_input_fn,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
113
                        custom_callbacks=None,
114
115
                        run_eagerly=False,
                        use_keras_compile_fit=False):
116
117
118
119
120
  """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():
121
    """Gets a classifier model."""
122
    classifier_model, core_model = (
123
124
125
126
        bert_models.classifier_model(
            bert_config,
            num_classes,
            max_seq_length,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
127
128
            hub_module_url=FLAGS.hub_module_url,
            hub_module_trainable=FLAGS.hub_module_trainable))
Hongkun Yu's avatar
Hongkun Yu committed
129
130
131
132
    optimizer = optimization.create_optimizer(initial_lr,
                                              steps_per_epoch * epochs,
                                              warmup_steps, FLAGS.end_lr,
                                              FLAGS.optimizer_type)
133
134
135
136
    classifier_model.optimizer = performance.configure_optimizer(
        optimizer,
        use_float16=common_flags.use_float16(),
        use_graph_rewrite=common_flags.use_graph_rewrite())
137
138
    return classifier_model, core_model

139
  loss_fn = get_loss_fn(num_classes)
140
141
142
143
144
145
146

  # 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)

147
  if use_keras_compile_fit:
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
148
149
    # 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
150
                 'distribution strategy.')
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
151
152
153
154
155
156
157
158
159
160
161
    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,
Hongkun Yu's avatar
Hongkun Yu committed
162
        steps_per_loop,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
163
        eval_steps,
164
        custom_callbacks=custom_callbacks)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
165
166
167

  # 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
168
               'distribution strategy.')
169
170
171
172
173
174
  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,
175
      steps_per_loop=steps_per_loop,
176
177
178
179
180
      epochs=epochs,
      train_input_fn=train_input_fn,
      eval_input_fn=eval_input_fn,
      eval_steps=eval_steps,
      init_checkpoint=init_checkpoint,
181
      sub_model_export_name=FLAGS.sub_model_export_name,
182
      metric_fn=metric_fn,
183
184
      custom_callbacks=custom_callbacks,
      run_eagerly=run_eagerly)
185
186


A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
187
188
189
190
191
192
193
194
195
196
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,
Hongkun Yu's avatar
Hongkun Yu committed
197
                          steps_per_loop,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
198
199
200
201
202
203
                          eval_steps,
                          custom_callbacks=None):
  """Runs BERT classifier model using Keras compile/fit API."""

  with strategy.scope():
    training_dataset = train_input_fn()
Le Hou's avatar
Le Hou committed
204
    evaluation_dataset = eval_input_fn() if eval_input_fn else None
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
205
206
207
208
209
210
211
    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()

Hongkun Yu's avatar
Hongkun Yu committed
212
213
214
215
216
    bert_model.compile(
        optimizer=optimizer,
        loss=loss_fn,
        metrics=[metric_fn()],
        experimental_steps_per_execution=steps_per_loop)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
217

218
219
    summary_dir = os.path.join(model_dir, 'summaries')
    summary_callback = tf.keras.callbacks.TensorBoard(summary_dir)
Hongkun Yu's avatar
Hongkun Yu committed
220
221
222
223
224
225
226
227
    checkpoint = tf.train.Checkpoint(model=bert_model, optimizer=optimizer)
    checkpoint_manager = tf.train.CheckpointManager(
        checkpoint,
        directory=model_dir,
        max_to_keep=None,
        step_counter=optimizer.iterations,
        checkpoint_interval=0)
    checkpoint_callback = keras_utils.SimpleCheckpoint(checkpoint_manager)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244

    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


Hongkun Yu's avatar
Hongkun Yu committed
245
246
247
248
def get_predictions_and_labels(strategy,
                               trained_model,
                               eval_input_fn,
                               return_probs=False):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
249
250
251
252
253
254
255
256
257
  """Obtains predictions of trained model on evaluation data.

  Note that list of labels is returned along with the predictions because the
  order changes on distributing dataset over TPU pods.

  Args:
    strategy: Distribution strategy.
    trained_model: Trained model with preloaded weights.
    eval_input_fn: Input function for evaluation data.
Hongkun Yu's avatar
Hongkun Yu committed
258
    return_probs: Whether to return probabilities of classes.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
259
260
261
262
263
264
265
266
267
268
269
270
271

  Returns:
    predictions: List of predictions.
    labels: List of gold labels corresponding to predictions.
  """

  @tf.function
  def test_step(iterator):
    """Computes predictions on distributed devices."""

    def _test_step_fn(inputs):
      """Replicated predictions."""
      inputs, labels = inputs
Hongkun Yu's avatar
Hongkun Yu committed
272
273
274
      logits = trained_model(inputs, training=False)
      probabilities = tf.nn.softmax(logits)
      return probabilities, labels
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
275

Hongkun Yu's avatar
Hongkun Yu committed
276
    outputs, labels = strategy.run(_test_step_fn, args=(next(iterator),))
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
277
278
279
280
281
282
283
284
285
    # outputs: current batch logits as a tuple of shard logits
    outputs = tf.nest.map_structure(strategy.experimental_local_results,
                                    outputs)
    labels = tf.nest.map_structure(strategy.experimental_local_results, labels)
    return outputs, labels

  def _run_evaluation(test_iterator):
    """Runs evaluation steps."""
    preds, golds = list(), list()
Hongkun Yu's avatar
Hongkun Yu committed
286
287
288
289
290
291
292
293
294
295
296
297
    try:
      with tf.experimental.async_scope():
        while True:
          probabilities, labels = test_step(test_iterator)
          for cur_probs, cur_labels in zip(probabilities, labels):
            if return_probs:
              preds.extend(cur_probs.numpy().tolist())
            else:
              preds.extend(tf.math.argmax(cur_probs, axis=1).numpy())
            golds.extend(cur_labels.numpy().tolist())
    except (StopIteration, tf.errors.OutOfRangeError):
      tf.experimental.async_clear_error()
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
298
299
300
301
302
303
304
305
306
    return preds, golds

  test_iter = iter(
      strategy.experimental_distribute_datasets_from_function(eval_input_fn))
  predictions, labels = _run_evaluation(test_iter)

  return predictions, labels


Hongkun Yu's avatar
Hongkun Yu committed
307
308
def export_classifier(model_export_path, input_meta_data, bert_config,
                      model_dir):
309
310
311
312
313
  """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.
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
314
315
316
    bert_config: Bert configuration file to define core bert layers.
    model_dir: The directory where the model weights and training/evaluation
      summaries are stored.
317
318
319
320
321
322

  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
323
324
  if not model_dir:
    raise ValueError('Export path is not specified: %s' % model_dir)
325

Zongwei Zhou's avatar
Zongwei Zhou committed
326
327
  # Export uses float32 for now, even if training uses mixed precision.
  tf.keras.mixed_precision.experimental.set_policy('float32')
328
  classifier_model = bert_models.classifier_model(
Hongkun Yu's avatar
Hongkun Yu committed
329
      bert_config, input_meta_data['num_labels'])[0]
330

331
  model_saving_utils.export_bert_model(
Hongkun Yu's avatar
Hongkun Yu committed
332
      model_export_path, model=classifier_model, checkpoint_dir=model_dir)
333
334


Hongkun Yu's avatar
Hongkun Yu committed
335
336
def run_bert(strategy,
             input_meta_data,
337
             model_config,
Hongkun Yu's avatar
Hongkun Yu committed
338
             train_input_fn=None,
Le Hou's avatar
Le Hou committed
339
             eval_input_fn=None,
340
341
             init_checkpoint=None,
             custom_callbacks=None):
342
  """Run BERT training."""
343
  # Enables XLA in Session Config. Should not be set for TPU.
344
  keras_utils.set_session_config(FLAGS.enable_xla)
345
  performance.set_mixed_precision_policy(common_flags.dtype())
346
347
348
349
350
351
352
353
354
355

  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
356

357
358
359
  if not custom_callbacks:
    custom_callbacks = []

360
  if FLAGS.log_steps:
Hongkun Yu's avatar
Hongkun Yu committed
361
362
363
364
365
    custom_callbacks.append(
        keras_utils.TimeHistory(
            batch_size=FLAGS.train_batch_size,
            log_steps=FLAGS.log_steps,
            logdir=FLAGS.model_dir))
366

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
367
  trained_model = run_bert_classifier(
368
      strategy,
369
      model_config,
370
371
372
373
      input_meta_data,
      FLAGS.model_dir,
      epochs,
      steps_per_epoch,
374
      FLAGS.steps_per_loop,
375
376
377
      eval_steps,
      warmup_steps,
      FLAGS.learning_rate,
Le Hou's avatar
Le Hou committed
378
      init_checkpoint or FLAGS.init_checkpoint,
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
379
380
      train_input_fn,
      eval_input_fn,
381
      run_eagerly=FLAGS.run_eagerly,
382
383
      use_keras_compile_fit=FLAGS.use_keras_compile_fit,
      custom_callbacks=custom_callbacks)
384

385
  if FLAGS.model_export_path:
386
    model_saving_utils.export_bert_model(
Hongkun Yu's avatar
Hongkun Yu committed
387
        FLAGS.model_export_path, model=trained_model)
388
389
  return trained_model

390

391
392
def custom_main(custom_callbacks=None):
  """Run classification.
393

394
395
396
  Args:
    custom_callbacks: list of tf.keras.Callbacks passed to training loop.
  """
Le Hou's avatar
Le Hou committed
397
398
  gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_param)

399
400
401
402
403
404
  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/'

Hongkun Yu's avatar
Hongkun Yu committed
405
406
407
408
409
410
411
  bert_config = bert_configs.BertConfig.from_json_file(FLAGS.bert_config_file)

  if FLAGS.mode == 'export_only':
    export_classifier(FLAGS.model_export_path, input_meta_data, bert_config,
                      FLAGS.model_dir)
    return

412
413
414
415
  strategy = distribution_utils.get_distribution_strategy(
      distribution_strategy=FLAGS.distribution_strategy,
      num_gpus=FLAGS.num_gpus,
      tpu_address=FLAGS.tpu)
Hongkun Yu's avatar
Hongkun Yu committed
416
  eval_input_fn = get_dataset_fn(
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
417
      FLAGS.eval_data_path,
Hongkun Yu's avatar
Hongkun Yu committed
418
      input_meta_data['max_seq_length'],
Hongkun Yu's avatar
Hongkun Yu committed
419
420
421
      FLAGS.eval_batch_size,
      is_training=False)

Hongkun Yu's avatar
Hongkun Yu committed
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
  if FLAGS.mode == 'predict':
    with strategy.scope():
      classifier_model = bert_models.classifier_model(
          bert_config, input_meta_data['num_labels'])[0]
      checkpoint = tf.train.Checkpoint(model=classifier_model)
      latest_checkpoint_file = (
          FLAGS.predict_checkpoint_path or
          tf.train.latest_checkpoint(FLAGS.model_dir))
      assert latest_checkpoint_file
      logging.info('Checkpoint file %s found and restoring from '
                   'checkpoint', latest_checkpoint_file)
      checkpoint.restore(
          latest_checkpoint_file).assert_existing_objects_matched()
      preds, _ = get_predictions_and_labels(
          strategy, classifier_model, eval_input_fn, return_probs=True)
    output_predict_file = os.path.join(FLAGS.model_dir, 'test_results.tsv')
    with tf.io.gfile.GFile(output_predict_file, 'w') as writer:
      logging.info('***** Predict results *****')
      for probabilities in preds:
        output_line = '\t'.join(
            str(class_probability)
            for class_probability in probabilities) + '\n'
        writer.write(output_line)
    return

  if FLAGS.mode != 'train_and_eval':
    raise ValueError('Unsupported mode is specified: %s' % FLAGS.mode)
  train_input_fn = get_dataset_fn(
      FLAGS.train_data_path,
      input_meta_data['max_seq_length'],
      FLAGS.train_batch_size,
      is_training=True)
  run_bert(
      strategy,
      input_meta_data,
      bert_config,
      train_input_fn,
      eval_input_fn,
      custom_callbacks=custom_callbacks)
461
462
463
464


def main(_):
  custom_main(custom_callbacks=None)
465
466
467
468
469


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