| import numpy as np |
|
|
| from mla.backend import xp |
| from mla.tensor import Tensor |
| from mla.optim import AdamW, clip_grad_norm |
|
|
|
|
| def test_adamw_converges_quadratic(): |
| xp.random.seed(0) |
| target = xp.asarray([3.0, -2.0, 0.5, 1.0]) |
| w = Tensor(xp.random.randn(4)) |
| opt = AdamW([w], lr=0.1, weight_decay=0.0) |
| loss = None |
| for _ in range(600): |
| opt.zero_grad() |
| diff = w.sub(target) |
| loss = diff.mul(diff).sum() |
| loss.backward() |
| opt.step() |
| assert float(loss.data) < 1e-4, float(loss.data) |
| assert np.allclose(np.asarray(w.data), np.asarray(target), atol=1e-2) |
|
|
|
|
| def test_adamw_first_step_scale(): |
| w = Tensor(xp.asarray([10.0])) |
| opt = AdamW([w], lr=0.1, weight_decay=0.0) |
| opt.zero_grad() |
| w.grad = xp.asarray([0.5]) |
| opt.step() |
| assert abs(float(w.data[0]) - 9.9) < 1e-6 |
|
|
|
|
| def test_weight_decay_pulls_to_zero(): |
| w = Tensor(xp.asarray([5.0])) |
| opt = AdamW([w], lr=0.1, weight_decay=0.1) |
| for _ in range(50): |
| opt.zero_grad() |
| w.grad = xp.zeros_like(w.data) |
| opt.step() |
| assert float(w.data[0]) < 5.0 |
|
|
|
|
| def test_clip_grad_norm_scales(): |
| p = Tensor(xp.asarray([3.0, 4.0])) |
| p.grad = xp.asarray([3.0, 4.0]) |
| total = clip_grad_norm([p], max_norm=1.0) |
| assert abs(total - 5.0) < 1e-6 |
| new_norm = float((np.asarray(p.grad) ** 2).sum()) ** 0.5 |
| assert abs(new_norm - 1.0) < 1e-4 |
|
|
|
|
| def test_clip_grad_norm_noop_under_threshold(): |
| p = Tensor(xp.asarray([0.3, 0.4])) |
| p.grad = xp.asarray([0.3, 0.4]) |
| total = clip_grad_norm([p], max_norm=1.0) |
| assert abs(total - 0.5) < 1e-6 |
| new_norm = float((np.asarray(p.grad) ** 2).sum()) ** 0.5 |
| assert abs(new_norm - 0.5) < 1e-6 |
|
|
