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

Function test_adam

imperative/python/test/integration/test_optimizer.py:172–203  ·  view source on GitHub ↗
(monkeypatch, case, update_lr, inplace_mode)

Source from the content-addressed store, hash-verified

170@pytest.mark.parametrize("update_lr", [False, True])
171@pytest.mark.parametrize("inplace_mode", [False, True])
172def test_adam(monkeypatch, case, update_lr, inplace_mode):
173 class CheckValue:
174 def __init__(self, net, **kwarg):
175 self.m_slots = {}
176 self.v_slots = {}
177 for param in net.parameters():
178 self.m_slots[param] = np.zeros(param.shape).astype(np.float32)
179 self.v_slots[param] = np.zeros(param.shape).astype(np.float32)
180 for k, v in kwarg.items():
181 setattr(self, k, v)
182
183 def __call__(self, ori_params, new_params, step):
184 for param in new_params:
185 grad = param.grad.numpy()
186 if hasattr(self, "weight_decay") and self.weight_decay != 0.0:
187 grad = grad + ori_params[param] * self.weight_decay
188 m = self.m_slots[param]
189 v = self.v_slots[param]
190 m *= self.betas[0]
191 m += (1 - self.betas[0]) * grad
192 v *= self.betas[1]
193 v += (1 - self.betas[1]) * grad * grad
194 delta = (m / (1 - self.betas[0] ** step)) / (
195 np.sqrt(v / (1 - self.betas[1] ** step)) + self.eps
196 )
197 np.testing.assert_almost_equal(
198 param.numpy(), ori_params[param] - self.lr * delta, decimal=6
199 )
200
201 with monkeypatch.context() as mk:
202 mk.setenv("MEGENGINE_INPLACE_UPDATE", str(int(inplace_mode)))
203 _test_optimizer("Adam", case, CheckValue, update_lr=update_lr)
204
205
206@pytest.mark.parametrize(

Callers

nothing calls this directly

Calls 3

strFunction · 0.85
_test_optimizerFunction · 0.85
contextMethod · 0.80

Tested by

no test coverage detected