run_pretraining.py 6.95 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

  return _bert_pretrain_loss_fn


def run_customized_training(strategy,
                            bert_config,
André Susano Pinto's avatar
André Susano Pinto committed
91
                            init_checkpoint,
92
93
94
95
                            max_seq_length,
                            max_predictions_per_seq,
                            model_dir,
                            steps_per_epoch,
96
                            steps_per_loop,
97
98
99
                            epochs,
                            initial_lr,
                            warmup_steps,
100
101
                            end_lr,
                            optimizer_type,
102
                            input_files,
103
                            train_batch_size,
Chen Chen's avatar
Chen Chen committed
104
105
                            use_next_sentence_label=True,
                            custom_callbacks=None):
106
107
  """Run BERT pretrain model training using low-level API."""

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

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

127
  trained_model = model_training_utils.run_customized_training_loop(
128
129
      strategy=strategy,
      model_fn=_get_pretrain_model,
130
131
      loss_fn=get_loss_fn(),
      scale_loss=FLAGS.scale_loss,
132
      model_dir=model_dir,
André Susano Pinto's avatar
André Susano Pinto committed
133
      init_checkpoint=init_checkpoint,
134
135
      train_input_fn=train_input_fn,
      steps_per_epoch=steps_per_epoch,
136
      steps_per_loop=steps_per_loop,
Chen Chen's avatar
Chen Chen committed
137
      epochs=epochs,
Chen Chen's avatar
Chen Chen committed
138
139
      sub_model_export_name='pretrained/bert_model',
      custom_callbacks=custom_callbacks)
140

141
142
  return trained_model

143

Chen Chen's avatar
Chen Chen committed
144
def run_bert_pretrain(strategy, custom_callbacks=None):
145
146
  """Runs BERT pre-training."""

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

  # Runs customized training loop.
Chen Chen's avatar
Chen Chen committed
152
  logging.info('Training using customized training loop TF 2.0 with distributed'
153
154
               'strategy.')

155
156
  performance.set_mixed_precision_policy(common_flags.dtype())

157
158
159
  return run_customized_training(
      strategy,
      bert_config,
André Susano Pinto's avatar
André Susano Pinto committed
160
      FLAGS.init_checkpoint,  # Used to initialize only the BERT submodel.
161
162
163
164
      FLAGS.max_seq_length,
      FLAGS.max_predictions_per_seq,
      FLAGS.model_dir,
      FLAGS.num_steps_per_epoch,
165
      FLAGS.steps_per_loop,
166
167
168
      FLAGS.num_train_epochs,
      FLAGS.learning_rate,
      FLAGS.warmup_steps,
169
170
      FLAGS.end_lr,
      FLAGS.optimizer_type,
171
      FLAGS.input_files,
172
      FLAGS.train_batch_size,
Chen Chen's avatar
Chen Chen committed
173
174
      FLAGS.use_next_sentence_label,
      custom_callbacks=custom_callbacks)
175
176
177
178


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

189
  run_bert_pretrain(strategy)
190
191
192
193


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