Logging Learning Rate

[wbrenton]

Is there a method for accessing the learning rate being used by the optimizer at each step during training from the schedule.py schedules?

[atgctg]

To access the learning rate during training you could:

1. Get the current step from the optimizer and pass it to the schedule

def fit(params: optax.Params, optimizer: optax.GradientTransformation) -> optax.Params:
  opt_state = optimizer.init(params)

  for i, (batch, labels) in enumerate(zip(TRAINING_DATA, LABELS)):
    params, opt_state, loss_value = step(params, opt_state, batch, labels)
    if i % 100 == 0:
      count = opt_state.inner_state[0].count # get current step
      lr = schedule(count) # get learning rate from schedule
      print(f'Step {i:3}, Loss: {loss_value:.3f}, Learning rate: {lr:.9f}')

  return params

params = fit(initial_params, optimizer)

Step   0, Loss: 5.624, Learning rate: 0.019999981
Step 100, Loss: 0.000, Learning rate: 0.992905796
Step 200, Loss: 0.000, Learning rate: 0.938947499
...
Step 900, Loss: 0.000, Learning rate: 0.026557088

2. Make the learning rate available using inject_hyperparams

# Wrap the optimizer to inject the hyperparameters
optimizer = optax.inject_hyperparams(optax.adamw)(learning_rate=schedule)

def fit(params: optax.Params, optimizer: optax.GradientTransformation) -> optax.Params:
  opt_state = optimizer.init(params)

  # Since we injected hyperparams, we can access them directly here
  print(f'Available hyperparams: {" ".join(opt_state.hyperparams.keys())}\n')

  for i, (batch, labels) in enumerate(zip(TRAINING_DATA, LABELS)):
    params, opt_state, loss_value = step(params, opt_state, batch, labels)
    if i % 100 == 0:
      # Get the updated learning rate
      lr = opt_state.hyperparams['learning_rate']
      print(f'Step {i:3}, Loss: {loss_value:.3f}, Learning rate: {lr:.3f}')

  return params

params = fit(initial_params, optimizer)

Step   0, Loss: 5.624, Learning rate: 0.020
Step 100, Loss: 0.000, Learning rate: 0.993
Step 200, Loss: 0.000, Learning rate: 0.939
...
Step 900, Loss: 0.000, Learning rate: 0.027

---

For further discussion, see #206