run_pretraining.py 6.64 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.
# ==============================================================================
Hongkun Yu's avatar
Hongkun Yu committed
15
"""Run masked LM/next sentence pre-training for BERT in TF 2.x."""
16
17
18
19
20
21
22
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl import app
from absl import flags
from absl import logging
Hongkun Yu's avatar
Hongkun Yu committed
23
import gin
24
import tensorflow as tf
25
from official.modeling import performance
26
from official.nlp import optimization
27
from official.nlp.bert import bert_models
28
from official.nlp.bert import common_flags
29
from official.nlp.bert import configs
30
from official.nlp.bert import input_pipeline
31
from official.nlp.bert import model_training_utils
32
from official.utils.misc import distribution_utils
33

34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

flags.DEFINE_string('input_files', None,
                    'File path to retrieve training data for pre-training.')
# Model training specific flags.
flags.DEFINE_integer(
    'max_seq_length', 128,
    'The maximum total input sequence length after WordPiece tokenization. '
    'Sequences longer than this will be truncated, and sequences shorter '
    'than this will be padded.')
flags.DEFINE_integer('max_predictions_per_seq', 20,
                     'Maximum predictions per sequence_output.')
flags.DEFINE_integer('train_batch_size', 32, 'Total batch size for training.')
flags.DEFINE_integer('num_steps_per_epoch', 1000,
                     'Total number of training steps to run per epoch.')
flags.DEFINE_float('warmup_steps', 10000,
                   'Warmup steps for Adam weight decay optimizer.')
50
51
flags.DEFINE_bool('use_next_sentence_label', True,
                  'Whether to use next sentence label to compute final loss.')
52

53
common_flags.define_common_bert_flags()
Hongkun Yu's avatar
Hongkun Yu committed
54
common_flags.define_gin_flags()
55

56
57
58
FLAGS = flags.FLAGS


Hongkun Yu's avatar
Hongkun Yu committed
59
def get_pretrain_dataset_fn(input_file_pattern, seq_length,
60
61
                            max_predictions_per_seq, global_batch_size,
                            use_next_sentence_label=True):
62
  """Returns input dataset from input file string."""
63
  def _dataset_fn(ctx=None):
64
    """Returns tf.data.Dataset for distributed BERT pretraining."""
Hongkun Yu's avatar
Hongkun Yu committed
65
    input_patterns = input_file_pattern.split(',')
Hongkun Yu's avatar
Hongkun Yu committed
66
    batch_size = ctx.get_per_replica_batch_size(global_batch_size)
67
    train_dataset = input_pipeline.create_pretrain_dataset(
Hongkun Yu's avatar
Hongkun Yu committed
68
        input_patterns,
69
70
71
72
        seq_length,
        max_predictions_per_seq,
        batch_size,
        is_training=True,
73
74
        input_pipeline_context=ctx,
        use_next_sentence_label=use_next_sentence_label)
75
76
    return train_dataset

Hongkun Yu's avatar
Hongkun Yu committed
77
  return _dataset_fn
78
79


80
def get_loss_fn():
81
82
83
  """Returns loss function for BERT pretraining."""

  def _bert_pretrain_loss_fn(unused_labels, losses, **unused_args):
84
    return tf.reduce_mean(losses)
85
86
87
88
89
90
91
92
93
94

  return _bert_pretrain_loss_fn


def run_customized_training(strategy,
                            bert_config,
                            max_seq_length,
                            max_predictions_per_seq,
                            model_dir,
                            steps_per_epoch,
95
                            steps_per_loop,
96
97
98
                            epochs,
                            initial_lr,
                            warmup_steps,
99
100
                            end_lr,
                            optimizer_type,
101
                            input_files,
102
103
                            train_batch_size,
                            use_next_sentence_label=True):
104
105
  """Run BERT pretrain model training using low-level API."""

Hongkun Yu's avatar
Hongkun Yu committed
106
107
  train_input_fn = get_pretrain_dataset_fn(input_files, max_seq_length,
                                           max_predictions_per_seq,
108
109
                                           train_batch_size,
                                           use_next_sentence_label)
110
111

  def _get_pretrain_model():
112
    """Gets a pretraining model."""
113
    pretrain_model, core_model = bert_models.pretrain_model(
114
115
        bert_config, max_seq_length, max_predictions_per_seq,
        use_next_sentence_label=use_next_sentence_label)
116
    optimizer = optimization.create_optimizer(
117
        initial_lr, steps_per_epoch * epochs, warmup_steps,
118
        end_lr, optimizer_type)
119
120
121
122
    pretrain_model.optimizer = performance.configure_optimizer(
        optimizer,
        use_float16=common_flags.use_float16(),
        use_graph_rewrite=common_flags.use_graph_rewrite())
123
124
    return pretrain_model, core_model

125
  trained_model = model_training_utils.run_customized_training_loop(
126
127
      strategy=strategy,
      model_fn=_get_pretrain_model,
128
129
      loss_fn=get_loss_fn(),
      scale_loss=FLAGS.scale_loss,
130
131
132
      model_dir=model_dir,
      train_input_fn=train_input_fn,
      steps_per_epoch=steps_per_epoch,
133
      steps_per_loop=steps_per_loop,
Chen Chen's avatar
Chen Chen committed
134
135
      epochs=epochs,
      sub_model_export_name='pretrained/bert_model')
136

137
138
  return trained_model

139
140
141
142

def run_bert_pretrain(strategy):
  """Runs BERT pre-training."""

143
  bert_config = configs.BertConfig.from_json_file(FLAGS.bert_config_file)
144
145
146
147
148
149
150
  if not strategy:
    raise ValueError('Distribution strategy is not specified.')

  # Runs customized training loop.
  logging.info('Training using customized training loop TF 2.0 with distrubuted'
               'strategy.')

151
152
  performance.set_mixed_precision_policy(common_flags.dtype())

153
154
155
156
157
158
159
  return run_customized_training(
      strategy,
      bert_config,
      FLAGS.max_seq_length,
      FLAGS.max_predictions_per_seq,
      FLAGS.model_dir,
      FLAGS.num_steps_per_epoch,
160
      FLAGS.steps_per_loop,
161
162
163
      FLAGS.num_train_epochs,
      FLAGS.learning_rate,
      FLAGS.warmup_steps,
164
165
      FLAGS.end_lr,
      FLAGS.optimizer_type,
166
      FLAGS.input_files,
167
168
      FLAGS.train_batch_size,
      FLAGS.use_next_sentence_label)
169
170
171
172


def main(_):
  # Users should always run this script under TF 2.x
Hongkun Yu's avatar
Hongkun Yu committed
173
  gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_param)
174
175
  if not FLAGS.model_dir:
    FLAGS.model_dir = '/tmp/bert20/'
176
177
178
179
  strategy = distribution_utils.get_distribution_strategy(
      distribution_strategy=FLAGS.distribution_strategy,
      num_gpus=FLAGS.num_gpus,
      tpu_address=FLAGS.tpu)
180
181
182
  if strategy:
    print('***** Number of cores used : ', strategy.num_replicas_in_sync)

183
  run_bert_pretrain(strategy)
184
185
186
187


if __name__ == '__main__':
  app.run(main)