(factors='constant * linear_warmup * rsqrt_decay',
                constant=0.1,  # pylint: disable=redefined-outer-name
                warmup_steps=400,
                decay_factor=0.5,
                steps_per_decay=20000,
                steps_per_cycle=100000,
                second_constant=0.01,
                second_constant_step=10000,
                minimum=0)
return learning_rate
Factor-based learning rate schedule.

  Interprets factors in the factors string which can consist of:
  * constant: interpreted as the constant value,
  * linear_warmup: interpreted as linear warmup until warmup_steps,
  * rsqrt_decay: divide by square root of max(step, warmup_steps)
  * decay_every: Every k steps decay the learning rate by decay_factor.
  * cosine_deay: Cyclic cosine decay, uses steps_per_cycle parameter.
  * two_constants: constant until second_constant_step, then switch to
    second_constant.

  Args:
    factors: a string with factors separated by '*' that defines the schedule.
    constant: float, the starting constant for the learning rate schedule.
    warmup_steps: how many steps to warm up for in the warmup schedule.
    decay_factor: The amount to decay the learning rate by.
    steps_per_decay: How often to decay the learning rate.
    steps_per_cycle: Steps per cycle when using cosine decay.
    second_constant: float, the second constant for the learning rate schedule.
    second_constant_step: the step when the second_constant is triggered.
    minimum: if the computed rate is below the minimum, then return the minimum.

  Returns:
    a function learning_rate(step): float -> {'learning_rate': float}, the
    step-dependent lr.
Factor-based learning rate schedule.
Factor - based learning rate schedule .
def multifactor(factors='constant * linear_warmup * rsqrt_decay',
                constant=0.1,  # pylint: disable=redefined-outer-name
                warmup_steps=400,
                decay_factor=0.5,
                steps_per_decay=20000,
                steps_per_cycle=100000,
                second_constant=0.01,
                second_constant_step=10000,
                minimum=0):
  """Factor-based learning rate schedule.

  Interprets factors in the factors string which can consist of:
  * constant: interpreted as the constant value,
  * linear_warmup: interpreted as linear warmup until warmup_steps,
  * rsqrt_decay: divide by square root of max(step, warmup_steps)
  * decay_every: Every k steps decay the learning rate by decay_factor.
  * cosine_deay: Cyclic cosine decay, uses steps_per_cycle parameter.
  * two_constants: constant until second_constant_step, then switch to
    second_constant.

  Args:
    factors: a string with factors separated by '*' that defines the schedule.
    constant: float, the starting constant for the learning rate schedule.
    warmup_steps: how many steps to warm up for in the warmup schedule.
    decay_factor: The amount to decay the learning rate by.
    steps_per_decay: How often to decay the learning rate.
    steps_per_cycle: Steps per cycle when using cosine decay.
    second_constant: float, the second constant for the learning rate schedule.
    second_constant_step: the step when the second_constant is triggered.
    minimum: if the computed rate is below the minimum, then return the minimum.

  Returns:
    a function learning_rate(step): float -> {'learning_rate': float}, the
    step-dependent lr.
  """
  factors = [n.strip() for n in factors.split('*')]

  def learning_rate(step):
    """Step to learning rate function."""
    ret = 1.0
    for name in factors:
      if name == 'constant':
        ret *= constant
      elif name == 'two_constants':
        if step < second_constant_step:
          ret *= constant
        else:
          ret *= second_constant
      elif name == 'linear_warmup':
        ret *= jnp.minimum(1.0, step / warmup_steps)
      elif name == 'rsqrt_decay':
        ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))
      elif name == 'rsqrt_normalized_decay':
        ret *= jnp.sqrt(warmup_steps)
        ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))
      elif name == 'decay_every':
        ret *= (decay_factor ** (step//steps_per_decay))
      elif name == 'cosine_decay':
        progress = jnp.maximum(
            0.0, (step - warmup_steps) / float(steps_per_cycle))
        ret *= (0.5 * (1.0 + jnp.cos(jnp.pi * (progress % 1.0))))
      else:
        raise ValueError('Unknown factor %s.' % name)
    # TODO(henrykm): return float(jnp.max(minimum, ret)) would be
    # better but causes TypeError: 'numpy.float64' object cannot
    # be interpreted as an integer
    if ret <= minimum:
      return minimum
    return ret

  return learning_rate
def multifactor ( factors = 'constant * linear_warmup * rsqrt_decay' , constant = 0.1 , # pylint: disable=redefined-outer-name warmup_steps = 400 , decay_factor = 0.5 , steps_per_decay = 20000 , steps_per_cycle = 100000 , second_constant = 0.01 , second_constant_step = 10000 , minimum = 0 ) : factors = [ n . strip ( ) for n in factors . split ( '*' ) ] def learning_rate ( step ) : """Step to learning rate function.""" ret = 1.0 for name in factors : if name == 'constant' : ret *= constant elif name == 'two_constants' : if step < second_constant_step : ret *= constant else : ret *= second_constant elif name == 'linear_warmup' : ret *= jnp . minimum ( 1.0 , step / warmup_steps ) elif name == 'rsqrt_decay' : ret /= jnp . sqrt ( jnp . maximum ( step , warmup_steps ) ) elif name == 'rsqrt_normalized_decay' : ret *= jnp . sqrt ( warmup_steps ) ret /= jnp . sqrt ( jnp . maximum ( step , warmup_steps ) ) elif name == 'decay_every' : ret *= ( decay_factor ** ( step // steps_per_decay ) ) elif name == 'cosine_decay' : progress = jnp . maximum ( 0.0 , ( step - warmup_steps ) / float ( steps_per_cycle ) ) ret *= ( 0.5 * ( 1.0 + jnp . cos ( jnp . pi * ( progress % 1.0 ) ) ) ) else : raise ValueError ( 'Unknown factor %s.' % name ) # TODO(henrykm): return float(jnp.max(minimum, ret)) would be # better but causes TypeError: 'numpy.float64' object cannot # be interpreted as an integer if ret <= minimum : return minimum return ret return learning_rate