zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-18 10:26:56 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Treate 'minimum' differently in cosine decay (#1138)

Browse Source
This commit is contained in:
jlwitthuhn
2024-05-20 08:00:48 -07:00
committed by GitHub
parent 0a558577bf
commit 7e5674d8be
2 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/mlx/optimizers/schedulers.py
View File

@@ -58,14 +58,14 @@ def step_decay(init: float, decay_rate: float, step_size: int) -> Callable:
return schedule
def cosine_decay(init: float, decay_steps: int, minimum: float = 0.0) -> Callable:
def cosine_decay(init: float, decay_steps: int, end: float = 0.0) -> Callable:
r"""Make a cosine decay scheduler.
Args:
init (float): Initial value.
decay_steps (int): Number of steps to decay over. The decayed
value is constant for steps beyond ``decay_steps``.
minimum (float, optional): Minimal value to decay to. Default: ``0``.
end (float, optional): Final value to decay to. Default: ``0``.
Example:
@@ -83,7 +83,7 @@ def cosine_decay(init: float, decay_steps: int, minimum: float = 0.0) -> Callabl
def scheduler(step):
s = mx.minimum(step, decay_steps)
decay = 0.5 * (1.0 + mx.cos((math.pi / decay_steps) * s))
return mx.maximum(init * decay, minimum)
return end + decay * (init - end)
return scheduler