CRG-optimizer-schedule.py

# Schedules with t from 0-1, eg use as lr_sch(t/steps)
def lr_sch(t):
  left_br = 20 * t - 5
  right_br = - (1.45 * t + 2.08)
  def denom(sign):
    return (1 + jnp.exp(- sign * (19 * (t - 0.015))))
  return 10 ** ((left_br / denom(-1)) + (right_br / denom(+1)))

def wd_sch(t):
  return 10 ** (-np.log(np.exp( 10.7 * t - 2.7) + 1) - 2 )

def b2_sch(t):
  return (1 - 10 ** (3.9 * np.exp( - 3.1 * (t - 0.17) ** 2) - 5.458))

def eps_sch(t):
  return 10 ** (-5.5 + 1.6 * np.exp(-67 * (t - 0.205) ** 2))