MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / test_lamb

Function test_lamb

imperative/python/test/unit/optimizer/test_lamb.py:67–87  ·  view source on GitHub ↗
()

Source from the content-addressed store, hash-verified

65
66@pytest.mark.skip(reason="pytest aborted, the same as groupnorm")
67def 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)

Callers

nothing calls this directly

Calls 6

LAMBUpdateClass · 0.85
lamb_updateFunction · 0.85
tensorMethod · 0.80
uniformMethod · 0.80
applyFunction · 0.50
numpyMethod · 0.45

Tested by

no test coverage detected