Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions TensorFlow/LanguageModeling/BERT/optimization.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
from horovod.tensorflow.compression import Compression

def create_optimizer(loss, init_lr, num_train_steps, num_warmup_steps, hvd=None, manual_fp16=False, use_fp16=False, num_accumulation_steps=1,
optimizer_type="adam", allreduce_post_accumulation=False):
optimizer_type="adam", allreduce_post_accumulation=False, init_loss_scale=2**32):
"""Creates an optimizer training op."""
global_step = tf.compat.v1.train.get_or_create_global_step()

Expand Down Expand Up @@ -96,11 +96,11 @@ def create_optimizer(loss, init_lr, num_train_steps, num_warmup_steps, hvd=None,
if hvd is not None and (num_accumulation_steps == 1 or (not allreduce_post_accumulation)):
optimizer = hvd.DistributedOptimizer(optimizer, sparse_as_dense=True, compression=Compression.fp16 if use_fp16 or manual_fp16 else Compression.none)
if use_fp16:
loss_scaler = tf.train.experimental.DynamicLossScale(initial_loss_scale=2**32, increment_period=1000, multiplier=2.0)
loss_scaler = tf.train.experimental.DynamicLossScale(initial_loss_scale=init_loss_scale, increment_period=1000, multiplier=2.0)
optimizer = tf.train.experimental.enable_mixed_precision_graph_rewrite(optimizer, loss_scaler)
loss_scale_value = tf.identity(loss_scaler(), name="loss_scale")
if manual_fp16:
loss_scale_manager = tf.contrib.mixed_precision.ExponentialUpdateLossScaleManager(init_loss_scale=2 ** 32,
loss_scale_manager = tf.contrib.mixed_precision.ExponentialUpdateLossScaleManager(init_loss_scale=init_loss_scale,
incr_every_n_steps=1000,
decr_every_n_nan_or_inf=2,
decr_ratio=0.5)
Expand Down Expand Up @@ -157,7 +157,7 @@ def update(accum_vars):
lambda: update(accum_vars), lambda: tf.no_op())

new_global_step = tf.cond(tf.math.logical_and(update_step,
tf.cast(hvd.allreduce(tf.cast(batch_finite, tf.int32)), tf.bool)) if hvd is not None else batch_finite,
tf.cast(hvd.allreduce(tf.cast(batch_finite, tf.int32)), tf.bool) if hvd is not None else batch_finite),
lambda: global_step+1,
lambda: global_step)
new_global_step = tf.identity(new_global_step, name='step_update')
Expand Down
2 changes: 1 addition & 1 deletion TensorFlow/LanguageModeling/BERT/run_classifier.py
Original file line number Diff line number Diff line change
Expand Up @@ -618,7 +618,7 @@ def main(_):
tf.compat.v1.logging.info("Total Inference Time = %0.2f for Sentences = %d", eval_time_elapsed,
eval_hooks[-1].count * FLAGS.eval_batch_size)
tf.compat.v1.logging.info("Total Inference Time W/O Overhead = %0.2f for Sentences = %d", eval_time_wo_overhead,
num_sentences))
num_sentences)
tf.compat.v1.logging.info("Summary Inference Statistics on EVAL set")
tf.compat.v1.logging.info("Batch size = %d", FLAGS.eval_batch_size)
tf.compat.v1.logging.info("Sequence Length = %d", FLAGS.max_seq_length)
Expand Down
28 changes: 17 additions & 11 deletions TensorFlow/LanguageModeling/BERT/run_pretraining.py
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,7 @@
flags.DEFINE_bool("use_xla", False, "Whether to enable XLA JIT compilation.")

flags.DEFINE_bool("use_fp16", False, "Whether to enable AMP ops.")
flags.DEFINE_integer("init_loss_scale", 2**32, "Initial value of loss scale if mixed precision training")

# report samples/sec, total loss and learning rate during training
class _LogSessionRunHook(tf.estimator.SessionRunHook):
Expand All @@ -142,7 +143,9 @@ def after_create_session(self, session, coord):
self.loss = 0.0 # accumulation of loss in each step between every print

self.total_time = 0.0 # total time taken to train (excluding warmup + ckpt saving steps)
self.step_time = 0.0 # time taken per step
self.init_global_step = session.run(tf.train.get_global_step()) # training starts at init_global_step
self.skipped = 0

def before_run(self, run_context):
self.t0 = time.time()
Expand Down Expand Up @@ -185,18 +188,25 @@ def after_run(self, run_context, run_values):
self.global_step, update_step, total_loss, lr, nsp_loss, mlm_loss = run_values.\
results

# Removing first two steps after every checkpoint save from timing
if (self.global_step - self.init_global_step) % self.save_ckpt_steps <= 5:
print("Skipping time record for ", self.global_step, " due to checkpoint-saving/warmup overhead")
else:
self.total_time += run_time
self.elapsed_secs += run_time
self.step_time += run_time

print_step = self.global_step + 1 # One-based index for printing.
self.loss += total_loss
self.all_count += 1
if update_step:

self.count += 1

# Removing first six steps after every checkpoint save from timing
if (self.global_step - self.init_global_step) % self.save_ckpt_steps < 6:
print("Skipping time record for ", self.global_step, " due to checkpoint-saving/warmup overhead")
self.skipped += 1
else:
self.total_time += self.step_time

self.step_time = 0.0 #Reset Step Time

if (print_step == 1 or print_step % self.display_every == 0):
dt = self.elapsed_secs / self.count
sent_per_sec = self.global_batch_size / dt
Expand Down Expand Up @@ -231,15 +241,11 @@ def after_run(self, run_context, run_values):
"total_loss":float(total_loss), "avg_loss_step":float(avg_loss_step),
"learning_rate": str(lr)},
verbosity=Verbosity.DEFAULT)

self.elapsed_secs = 0.0
self.count = 0
self.loss = 0.0
self.all_count = 0
def end(self, session):
num_global_steps = self.global_step - self.init_global_step
self.skipped = (num_global_steps // self.save_ckpt_steps) * 5 + \
min(5, num_global_steps % self.save_ckpt_steps)


def model_fn_builder(bert_config, init_checkpoint, learning_rate,
num_train_steps, num_warmup_steps,
Expand Down Expand Up @@ -310,7 +316,7 @@ def model_fn(features, labels, mode, params): # pylint: disable=unused-argument
if mode == tf.estimator.ModeKeys.TRAIN:
train_op = optimization.create_optimizer(
total_loss, learning_rate, num_train_steps, num_warmup_steps,
hvd, FLAGS.manual_fp16, FLAGS.use_fp16, FLAGS.num_accumulation_steps, FLAGS.optimizer_type, FLAGS.allreduce_post_accumulation)
hvd, FLAGS.manual_fp16, FLAGS.use_fp16, FLAGS.num_accumulation_steps, FLAGS.optimizer_type, FLAGS.allreduce_post_accumulation, FLAGS.init_loss_scale)

output_spec = tf.estimator.EstimatorSpec(
mode=mode,
Expand Down
6 changes: 5 additions & 1 deletion TensorFlow/LanguageModeling/BERT/run_squad.py
Original file line number Diff line number Diff line change
Expand Up @@ -121,8 +121,10 @@ def extract_run_squad_flags():
"Proportion of training to perform linear learning rate warmup for. "
"E.g., 0.1 = 10% of training.")

flags.DEFINE_integer("save_checkpoints_steps", 1000,
flags.DEFINE_integer("save_checkpoints_steps", 5000,
"How often to save the model checkpoint.")
flags.DEFINE_integer("display_loss_steps", 10,
"How often to print loss from estimator")

flags.DEFINE_integer("iterations_per_loop", 1000,
"How many steps to make in each estimator call.")
Expand Down Expand Up @@ -967,6 +969,8 @@ def main(_):
model_dir=FLAGS.output_dir if master_process else None,
session_config=config,
save_checkpoints_steps=FLAGS.save_checkpoints_steps if master_process else None,
save_summary_steps=FLAGS.save_checkpoints_steps if master_process else None,
log_step_count_steps=FLAGS.display_loss_steps,
keep_checkpoint_max=1)

if master_process:
Expand Down