(loss, init_lr, num_train_steps, num_warmup_steps, use_tpu)
return train_op
Creates an optimizer training op.
Creates an optimizer training op.
Creates an optimizer training op .
def create_optimizer(loss, init_lr, num_train_steps, num_warmup_steps, use_tpu):
  """Creates an optimizer training op."""
  global_step = tf.train.get_or_create_global_step()

  learning_rate = tf.constant(value=init_lr, shape=[], dtype=tf.float32)

  # Implements linear decay of the learning rate.
  learning_rate = tf.train.polynomial_decay(
      learning_rate,
      global_step,
      num_train_steps,
      end_learning_rate=0.0,
      power=1.0,
      cycle=False)

  # Implements linear warmup. I.e., if global_step < num_warmup_steps, the
  # learning rate will be `global_step/num_warmup_steps * init_lr`.
  if num_warmup_steps:
    global_steps_int = tf.cast(global_step, tf.int32)
    warmup_steps_int = tf.constant(num_warmup_steps, dtype=tf.int32)

    global_steps_float = tf.cast(global_steps_int, tf.float32)
    warmup_steps_float = tf.cast(warmup_steps_int, tf.float32)

    warmup_percent_done = global_steps_float / warmup_steps_float
    warmup_learning_rate = init_lr * warmup_percent_done

    is_warmup = tf.cast(global_steps_int < warmup_steps_int, tf.float32)
    learning_rate = (
        (1.0 - is_warmup) * learning_rate + is_warmup * warmup_learning_rate)

  # It is recommended that you use this optimizer for fine tuning, since this
  # is how the model was trained (note that the Adam m/v variables are NOT
  # loaded from init_checkpoint.)
  optimizer = AdamWeightDecayOptimizer(
      learning_rate=learning_rate,
      weight_decay_rate=0.01,
      beta_1=0.9,
      beta_2=0.999, # 0.98 ONLY USED FOR PRETRAIN. MUST CHANGE AT FINE-TUNING 0.999,
      epsilon=1e-6,
      exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"])

  if use_tpu:
    optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)

  tvars = tf.trainable_variables()
  grads = tf.gradients(loss, tvars)

  # This is how the model was pre-trained.
  (grads, _) = tf.clip_by_global_norm(grads, clip_norm=1.0)

  train_op = optimizer.apply_gradients(
      zip(grads, tvars), global_step=global_step)

  # Normally the global step update is done inside of `apply_gradients`.
  # However, `AdamWeightDecayOptimizer` doesn't do this. But if you use
  # a different optimizer, you should probably take this line out.
  new_global_step = global_step + 1
  train_op = tf.group(train_op, [global_step.assign(new_global_step)])
  return train_op
def create_optimizer ( loss , init_lr , num_train_steps , num_warmup_steps , use_tpu ) : global_step = tf . train . get_or_create_global_step ( ) learning_rate = tf . constant ( value = init_lr , shape = [ ] , dtype = tf . float32 ) # Implements linear decay of the learning rate. learning_rate = tf . train . polynomial_decay ( learning_rate , global_step , num_train_steps , end_learning_rate = 0.0 , power = 1.0 , cycle = False ) # Implements linear warmup. I.e., if global_step < num_warmup_steps, the # learning rate will be `global_step/num_warmup_steps * init_lr`. if num_warmup_steps : global_steps_int = tf . cast ( global_step , tf . int32 ) warmup_steps_int = tf . constant ( num_warmup_steps , dtype = tf . int32 ) global_steps_float = tf . cast ( global_steps_int , tf . float32 ) warmup_steps_float = tf . cast ( warmup_steps_int , tf . float32 ) warmup_percent_done = global_steps_float / warmup_steps_float warmup_learning_rate = init_lr * warmup_percent_done is_warmup = tf . cast ( global_steps_int < warmup_steps_int , tf . float32 ) learning_rate = ( ( 1.0 - is_warmup ) * learning_rate + is_warmup * warmup_learning_rate ) # It is recommended that you use this optimizer for fine tuning, since this # is how the model was trained (note that the Adam m/v variables are NOT # loaded from init_checkpoint.) optimizer = AdamWeightDecayOptimizer ( learning_rate = learning_rate , weight_decay_rate = 0.01 , beta_1 = 0.9 , beta_2 = 0.999 , # 0.98 ONLY USED FOR PRETRAIN. MUST CHANGE AT FINE-TUNING 0.999, epsilon = 1e-6 , exclude_from_weight_decay = [ "LayerNorm" , "layer_norm" , "bias" ] ) if use_tpu : optimizer = tf . contrib . tpu . CrossShardOptimizer ( optimizer ) tvars = tf . trainable_variables ( ) grads = tf . gradients ( loss , tvars ) # This is how the model was pre-trained. ( grads , _ ) = tf . clip_by_global_norm ( grads , clip_norm = 1.0 ) train_op = optimizer . apply_gradients ( zip ( grads , tvars ) , global_step = global_step ) # Normally the global step update is done inside of `apply_gradients`. # However, `AdamWeightDecayOptimizer` doesn't do this. But if you use # a different optimizer, you should probably take this line out. new_global_step = global_step + 1 train_op = tf . group ( train_op , [ global_step . assign ( new_global_step ) ] ) return train_op