()
| 65 | |
| 66 | @pytest.mark.skip(reason="pytest aborted, the same as groupnorm") |
| 67 | def test_lamb(): |
| 68 | op = LAMBUpdate(0.9, 0.999, 1, 1e-3, 0.4, 1e-8, True, False) |
| 69 | m_t_1 = mge.tensor(np.random.uniform(size=(256, 256)), dtype=np.float32) |
| 70 | v_t_1 = mge.tensor(np.random.uniform(size=(256, 256)), dtype=np.float32) |
| 71 | params = mge.tensor(np.random.uniform(size=(256, 256)), dtype=np.float32) |
| 72 | grad = mge.tensor(np.random.uniform(size=(256, 256)), dtype=np.float16) |
| 73 | (new_m_t, new_v_t, new_param) = apply(op, m_t_1, v_t_1, params, grad) |
| 74 | |
| 75 | param_group = { |
| 76 | "betas": (0.9, 0.999), |
| 77 | "step": 1, |
| 78 | "lr": 1e-3, |
| 79 | "weight_decay": 0.4, |
| 80 | "eps": 1e-8, |
| 81 | } |
| 82 | gt_m_t, gt_v_t, gt_new_param = lamb_update( |
| 83 | param_group, 1, m_t_1, v_t_1, params, grad, True, False |
| 84 | ) |
| 85 | np.testing.assert_allclose(new_m_t.numpy(), gt_m_t.numpy(), atol=1e-2) |
| 86 | np.testing.assert_allclose(new_v_t.numpy(), gt_v_t.numpy(), atol=1e-2) |
| 87 | np.testing.assert_allclose(new_param.numpy(), gt_new_param.numpy(), atol=1e-2) |
nothing calls this directly
no test coverage detected