From 376e65a67481bcd370c77b119773b11bb612b0c3 Mon Sep 17 00:00:00 2001 From: Timothy Liu Date: Sun, 13 Oct 2019 11:04:49 +0000 Subject: [PATCH] Added automatic mixed precision and XLA options to run_tf_glue.py --- examples/run_tf_glue.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/examples/run_tf_glue.py b/examples/run_tf_glue.py index c05420d680..399fe9e616 100644 --- a/examples/run_tf_glue.py +++ b/examples/run_tf_glue.py @@ -6,6 +6,11 @@ from transformers import BertTokenizer, TFBertForSequenceClassification, glue_co # script parameters BATCH_SIZE = 32 EVAL_BATCH_SIZE = BATCH_SIZE * 2 +USE_XLA = False +USE_AMP = False + +tf.config.optimizer.set_jit(USE_XLA) +tf.config.optimizer.set_experimental_options({"auto_mixed_precision": USE_AMP}) # Load tokenizer and model from pretrained model/vocabulary tokenizer = BertTokenizer.from_pretrained('bert-base-cased') @@ -23,10 +28,13 @@ train_dataset = train_dataset.shuffle(128).batch(BATCH_SIZE).repeat(-1) valid_dataset = valid_dataset.batch(EVAL_BATCH_SIZE) # Prepare training: Compile tf.keras model with optimizer, loss and learning rate schedule -optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08) +opt = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08) +if USE_AMP: + # loss scaling is currently required when using mixed precision + opt = tf.keras.mixed_precision.experimental.LossScaleOptimizer(opt, 'dynamic') loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) metric = tf.keras.metrics.SparseCategoricalAccuracy('accuracy') -model.compile(optimizer=optimizer, loss=loss, metrics=[metric]) +model.compile(optimizer=opt, loss=loss, metrics=[metric]) # Train and evaluate using tf.keras.Model.fit() train_steps = train_examples//BATCH_SIZE