From 80b2296bdb7ed60f219d1ce04017f17632dd1ef2 Mon Sep 17 00:00:00 2001 From: Jason Wang Date: Wed, 20 Nov 2024 11:11:19 -0800 Subject: [PATCH] Update test_optimizer_config.py (#814) Add simple tests for rewarmup schedule --- tests/test_optimizer_config.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_optimizer_config.py b/tests/test_optimizer_config.py index f8227e43d..9c5b91d7c 100644 --- a/tests/test_optimizer_config.py +++ b/tests/test_optimizer_config.py @@ -172,6 +172,7 @@ def test_rewarmup_schedule(): assert np.isclose(sched_fn(100), 1e-2) # Warmup reaches max LR # First decay phase + assert np.isclose(sched_fn(300), 0.6e-2) # Mid of first decay assert np.isclose(sched_fn(500), 0.2e-2) # End of first decay # Rewarmup at start of second cycle @@ -184,4 +185,5 @@ def test_rewarmup_schedule(): assert sched_fn(rewarmup_end + 1) < sched_fn(rewarmup_end) # Final decay phase + assert sched_fn(999 - 1) > sched_fn(999) assert np.isclose(sched_fn(999), 0.2e-2, atol=1e-4) # End of second decay