Inconsistencies in schedules API

[fabianp]

• For most schedules, the total number of steps is specified through the transition_steps parameter, but in some cases (e.g., optax.cosine_decay_schedule, optax.warmup_cosine_decay_schedule but confusingly not optax.cosine_onecycle_schedule) it's called decay_steps instead.
• The name sgdr_schedule is not descriptive of what the schedule actually does.
• Most warm-up learning rates like linear_onecycle_schedule and cosine_onecycle_schedule specify the length of the warm-up phrase using parameter pct_start , but warmup_cosine_decay_schedule instead specifies it through a parameter warmup_steps

---

In the documentation:
5. In the API reference https://optax.readthedocs.io/en/latest/api/optimizer_schedules.html there's a section "Schedules with warm-up". I would consider optax.cosine_onecycle_schedule to have warm-up, yet it's not in this section. My recommendation would be to remove the section ""Schedules with warm-up" and put optax.warmup_cosine_decay_schedule in the Cosine decay schedule section and optax.warmup_exponential_decay_schedule in the exponential decay section

[Abhinavcode13]

I'm ready to pick this up! @fabianp

[fabianp]

excellent @Abhinavcode13 ! which task would you like to pick up? I would suggest to pick one of the remaining 3 and start with that. Smaller PR are better for everyone :-)

[Abhinavcode13]

Sure

[Abhinavcode13]

• For most schedules, the total number of steps is specified through the transition_steps parameter, but in some cases (e.g., optax.cosine_decay_schedule, optax.warmup_cosine_decay_schedule but confusingly not optax.cosine_onecycle_schedule) it's called decay_steps instead.
• The name sgdr_schedule is not descriptive of what the schedule actually does.
• Most warm-up learning rates like linear_onecycle_schedule and cosine_onecycle_schedule specify the length of the warm-up phrase using parameter pct_start , but warmup_cosine_decay_schedule instead specifies it through a parameter warmup_steps

In the documentation: 5. In the API reference https://optax.readthedocs.io/en/latest/api/optimizer_schedules.html there's a section "Schedules with warm-up". I would consider optax.cosine_onecycle_schedule to have warm-up, yet it's not in this section. My recommendation would be to remove the section ""Schedules with warm-up" and put optax.warmup_cosine_decay_schedule in the Cosine decay schedule section and optax.warmup_exponential_decay_schedule in the exponential decay section

FYI: I would look up the second one first.

[bhargavyagnik]

I can take this up, If @Abhinavcode13 is busy. I have some questions:

• We need to ensure backwards compatibility too, right? I mean after changing transition_steps for cosine_decay and other functions, we need to keep decay_steps too right? Or raise an error/warning to update the change?
• While reviewing the codebase, I observed that for warmup_exponential_decay_schedule, the exponential decay has transition_steps = transition_steps and while warmup_cosine_decay_schedule has transition_steps = transition_steps - warmup_steps. Just wanted to ensure that transition steps passed after warmup should be as current or is there something wrong in any one?

The warmup_exponential_decay_schedule has a structure like :

 schedules = [
      linear_schedule(
          init_value=init_value,
          end_value=peak_value,
          transition_steps=warmup_steps,
      ),
      exponential_decay(
          init_value=peak_value,
          transition_steps=transition_steps,
          decay_rate=decay_rate,
          transition_begin=transition_begin,
          staircase=staircase,
          end_value=end_value,
      )

While warmup_cosine_decay_schedule would go from

  schedules = [
      linear_schedule(
          init_value=init_value,
          end_value=peak_value,
          transition_steps=warmup_steps,
      ),
      cosine_decay_schedule(
          init_value=peak_value,
          transition_steps=transition_steps - warmup_steps,
          alpha=alpha,
          exponent=exponent,
      )