Optimizer, schedule and clipping tests
Test file
A test file. It is not part of the pipeline: neither train.main nor talk.main reaches it. It is listed with what it verifies on Checkpoints & Tests.
optim.pynumpy| Concept | What this file does | Lines | Status |
|---|---|---|---|
| AdamW update | Tests it | 13, 29 | Implemented |
| Learning-rate schedule | Tests it | 35 | Implemented |
| Gradient clipping | Tests it | 41, 49 | Implemented |
def test_adam_two_step_hand_trace()testLines 6-13GitHub: tests.test_optim.test_adam_two_step_hand_traceTwo steps match hand values
optim.Adam, optim.Adam.__init__, optim.Adam.stepdef test_adam_first_step_is_lr_times_sign_of_gradient()testLines 16-21GitHub: tests.test_optim.test_adam_first_step_is_lr_times_sign_of_gradientFirst update is -lr * sign(g)
optim.Adam, optim.Adam.__init__, optim.Adam.stepdef test_weight_decay_only_touches_matrices()testLines 24-29GitHub: tests.test_optim.test_weight_decay_only_touches_matricesDecay changes the 2-D array, not the 1-D array
optim.Adam, optim.Adam.__init__, optim.Adam.stepSchedule values match
optim.lr_scheduledef test_clip_rescales_all_tensors_together()testLines 38-43GitHub: tests.test_optim.test_clip_rescales_all_tensors_togetherGlobal clipping; returns the norm before clipping
optim.clip_grad_normdef test_clip_leaves_small_gradients_alone()testLines 46-49GitHub: tests.test_optim.test_clip_leaves_small_gradients_aloneSmall gradients are unchanged
optim.clip_grad_normimport numpy as npfrom optim import Adam, lr_schedule, clip_grad_normdef test_adam_two_step_hand_trace(): # theta0 = 1, lr = 0.1, g = 0.5 then 0.25 params = {"w": np.array([1.0])} opt = Adam(params) opt.step(params, {"w": np.array([0.5])}, lr=0.1) assert np.isclose(params["w"][0], 0.9000, atol=1e-4) opt.step(params, {"w": np.array([0.25])}, lr=0.1) assert np.isclose(params["w"][0], 0.8068, atol=1e-4)def test_adam_first_step_is_lr_times_sign_of_gradient(): # bias correction makes mhat/sqrt(vhat) = g/|g| at t = 1, whatever the size of g for g in (1e-4, 3.0, -250.0): params = {"w": np.array([0.0])} Adam(params).step(params, {"w": np.array([g])}, lr=0.01) assert np.isclose(params["w"][0], -0.01 * np.sign(g), atol=1e-6)def test_weight_decay_only_touches_matrices(): params = {"W": np.ones((2, 2)), "b": np.ones(2)} zeros = {k: np.zeros_like(v) for k, v in params.items()} Adam(params, weight_decay=0.1).step(params, zeros, lr=0.5) assert np.allclose(params["W"], 1 - 0.5 * 0.1) # shrunk assert np.allclose(params["b"], 1.0) # untoucheddef test_lr_schedule_values(): # W = 10, Ttot = 100, lr_max = 1e-3, lr_min = 1e-4, step is 1-based got = [lr_schedule(t, 10, 100, 1e-3, 1e-4) for t in (1, 5, 10, 55, 100)] assert np.allclose(got, [1e-4, 5e-4, 1e-3, 5.5e-4, 1e-4])def test_clip_rescales_all_tensors_together(): grads = {"a": np.array([3.0, 4.0]), "b": np.array([0.0])} norm = clip_grad_norm(grads, 1.0) assert np.isclose(norm, 5.0) # norm BEFORE clipping assert np.allclose(grads["a"], [0.6, 0.8]) assert np.isclose(np.sqrt(sum((g ** 2).sum() for g in grads.values())), 1.0)def test_clip_leaves_small_gradients_alone(): grads = {"a": np.array([0.3, 0.4])} clip_grad_norm(grads, 1.0) assert np.allclose(grads["a"], [0.3, 0.4])