(monkeypatch, case, update_lr, inplace_mode)
| 305 | @pytest.mark.parametrize("update_lr", [False, True]) |
| 306 | @pytest.mark.parametrize("inplace_mode", [False, True]) |
| 307 | def test_adamw(monkeypatch, case, update_lr, inplace_mode): |
| 308 | class CheckValue: |
| 309 | def __init__(self, net, **kwarg): |
| 310 | self.m_slots = {} |
| 311 | self.v_slots = {} |
| 312 | for param in net.parameters(): |
| 313 | self.m_slots[param] = np.zeros(param.shape).astype(np.float32) |
| 314 | self.v_slots[param] = np.zeros(param.shape).astype(np.float32) |
| 315 | self.weight_decay = 0.01 |
| 316 | for k, v in kwarg.items(): |
| 317 | setattr(self, k, v) |
| 318 | |
| 319 | def __call__(self, ori_params, new_params, step): |
| 320 | step = np.array(step).astype(np.float32) |
| 321 | for param in new_params: |
| 322 | grad = param.grad.numpy() |
| 323 | m = self.m_slots[param] |
| 324 | v = self.v_slots[param] |
| 325 | m *= self.betas[0] |
| 326 | m += (1 - self.betas[0]) * grad |
| 327 | v *= self.betas[1] |
| 328 | v += (1 - self.betas[1]) * grad * grad |
| 329 | delta = (m / (1 - self.betas[0] ** step)) / ( |
| 330 | np.sqrt(v / (1 - self.betas[1] ** step)) + self.eps |
| 331 | ) |
| 332 | delta += ori_params[param] * self.weight_decay |
| 333 | np.testing.assert_almost_equal( |
| 334 | param.numpy(), ori_params[param] - self.lr * delta, decimal=6 |
| 335 | ) |
| 336 | |
| 337 | with monkeypatch.context() as mk: |
| 338 | mk.setenv("MEGENGINE_INPLACE_UPDATE", str(int(inplace_mode))) |
| 339 | _test_optimizer("AdamW", case, CheckValue, update_lr=update_lr) |
nothing calls this directly
no test coverage detected