def _get_pretrain_model():
"""Gets a pretraining model."""
pretrain_model, core_model = bert_models.pretrain_model(
bert_config, max_seq_length, max_predictions_per_seq)
optimizer = optimization.create_optimizer(
initial_lr, steps_per_epoch * epochs, warmup_steps)
pretrain_model.optimizer = performance.configure_optimizer(
optimizer,
use_float16=common_flags.use_float16(),
use_graph_rewrite=common_flags.use_graph_rewrite())
return pretrain_model, core_model
def _get_pretrained_model(end_lr=0.0):
"""Gets a pretraining model."""
pretrain_model, core_model = bert_models.pretrain_model(
model_config, args.max_seq_length, args.max_predictions_per_seq)
optimizer = utils.optimizer.create_optimizer(
args.learning_rate, steps_per_epoch * args.num_epochs,
warmup_steps, args.end_lr, args.optimizer_type)
pretrain_model.optimizer = configure_optimizer(
optimizer,
use_float16=args.dtype == 'fp16',
use_graph_rewrite=False)
return pretrain_model, core_model
def _get_pretrain_model():
"""Gets a pretraining model."""
pretrain_model, core_model = bert_models.pretrain_model(
bert_config,
max_seq_length,
max_predictions_per_seq,
use_next_sentence_label=use_next_sentence_label)
optimizer = optimization.create_optimizer(initial_lr,
steps_per_epoch * epochs,
warmup_steps, end_lr,
optimizer_type)
pretrain_model.optimizer = performance.configure_optimizer(
optimizer, use_float16=common_flags.use_float16())
return pretrain_model, core_model
def _get_pretrain_model():
"""Gets a pretraining model."""
pretrain_model, core_model = bert_models.pretrain_model(
bert_config, max_seq_length, max_predictions_per_seq)
pretrain_model.optimizer = optimization.create_optimizer(
initial_lr, steps_per_epoch * epochs, warmup_steps)
if FLAGS.fp16_implementation == 'graph_rewrite':
# Note: when flags_obj.fp16_implementation == "graph_rewrite", dtype as
# determined by flags_core.get_tf_dtype(flags_obj) would be 'float32'
# which will ensure tf.compat.v2.keras.mixed_precision and
# tf.train.experimental.enable_mixed_precision_graph_rewrite do not double
# up.
pretrain_model.optimizer = tf.train.experimental.enable_mixed_precision_graph_rewrite(
pretrain_model.optimizer)
return pretrain_model, core_model
def test_pretrain_model(self):
model, encoder = bert_models.pretrain_model(
self._bert_test_config,
seq_length=5,
max_predictions_per_seq=2,
initializer=None,
use_next_sentence_label=True)
self.assertIsInstance(model, tf.keras.Model)
self.assertIsInstance(encoder, networks.TransformerEncoder)
# model has one scalar output: loss value.
self.assertEqual(model.output.shape.as_list(), [None,])
# Expect two output from encoder: sequence and classification output.
self.assertIsInstance(encoder.output, list)
self.assertLen(encoder.output, 2)
def _get_pretrained_model(end_lr=0.0):
"""Gets a pretraining model."""
pretrain_model, core_model = bert_models.pretrain_model(model_config, args.max_seq_length, args.max_predictions_per_seq)
if args.warmup_proportion is None:
warmup_steps = args.warmup_steps
warmup_proportion_perc = 100 * args.warmup_steps/(args.num_epochs * args.num_steps_per_epoch)
else:
warmup_steps = int(args.num_epochs * args.num_steps_per_epoch * args.warmup_proportion)
warmup_proportion_perc = args.warmup_proportion * 100
logger.info(f'Running {warmup_steps:,} warmup steps ({warmup_proportion_perc:.2f}% warmup)')
optimizer = utils.optimizer.create_optimizer(
args.learning_rate,
args.num_steps_per_epoch * args.num_epochs,
warmup_steps,
args.end_lr,
args.optimizer_type)
pretrain_model.optimizer = configure_optimizer(optimizer, use_float16=args.dtype == 'fp16', use_graph_rewrite=False)
return pretrain_model, core_model
