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

Function test_sgd

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

Source from the content-addressed store, hash-verified

125@pytest.mark.parametrize("update_lr", [False, True])
126@pytest.mark.parametrize("inplace_mode", [False, True])
127def test_sgd(monkeypatch, case, update_lr, inplace_mode):
128 class CheckValue:
129 def __init__(self, net, **kwarg):
130 self.slots = {}
131 for param in net.parameters():
132 self.slots[param] = np.zeros(param.shape).astype(np.float32)
133 for k, v in kwarg.items():
134 setattr(self, k, v)
135
136 def __call__(self, ori_params, new_params, step):
137 for param in new_params:
138 grad = param.grad.numpy()
139 if hasattr(self, "weight_decay") and self.weight_decay != 0.0:
140 grad = grad + ori_params[param] * self.weight_decay
141 if hasattr(self, "momentum") and self.momentum != 0.0:
142 self.slots[param] = grad + self.slots[param] * self.momentum
143 if hasattr(self, "nesterov") and self.nesterov:
144 delta = -self.lr * (grad + self.slots[param] * self.momentum)
145 else:
146 delta = -self.lr * self.slots[param]
147 else:
148 delta = -self.lr * grad
149 np.testing.assert_almost_equal(
150 param.numpy(), ori_params[param] + delta, decimal=6
151 )
152
153 with monkeypatch.context() as mk:
154 mk.setenv("MEGENGINE_INPLACE_UPDATE", str(int(inplace_mode)))
155 _test_optimizer("SGD", case, CheckValue, update_lr=update_lr)
156
157
158@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