File size: 993 Bytes
2ea2835 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 | from types import SimpleNamespace
import pytest
from trainer.distillation import _cosine_learning_rate, _learning_rate_schedule_identity
def test_step80_keeps_the_original_cosine_curve():
cfg = SimpleNamespace(lr_schedule="cosine", lr=2.5e-6, lr_final=5e-7,
lr_schedule_total_steps=200, max_steps=80)
early = _learning_rate_schedule_identity(cfg)
cfg.max_steps = 200
original = _learning_rate_schedule_identity(cfg)
assert early == original
for update in range(80):
assert _cosine_learning_rate(update, start=early["start"], end=early["end"], total_updates=early["total_updates"]) == _cosine_learning_rate(update, start=2.5e-6, end=5e-7, total_updates=200)
assert _cosine_learning_rate(79, start=2.5e-6, end=5e-7, total_updates=200) > 5e-7
cfg.max_steps = 1
assert _learning_rate_schedule_identity(cfg) == original
cfg.max_steps = 201
with pytest.raises(ValueError):
_learning_rate_schedule_identity(cfg)
|