From 66b6daf7777da27c7328a34cc34f945439862e6b Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Tue, 8 Sep 2020 17:21:02 +0800 Subject: [PATCH] test(mge/optimizer): fix test for new optimizer api GitOrigin-RevId: 482ee6265224f7fcc4c21300a597d36b4333cca3 --- .../python/megengine/distributed/functional.py | 2 +- .../test/integration/test_dp_correctness.py | 2 +- .../python/test/integration/test_optimizer.py | 11 ++++++----- .../python/test/integration/test_sgd_momentum.py | 15 ++++++++------- .../python/test/integration/test_trace_dump.py | 8 +++++--- 5 files changed, 21 insertions(+), 17 deletions(-) diff --git a/imperative/python/megengine/distributed/functional.py b/imperative/python/megengine/distributed/functional.py index 097ecb49..2bb01ccb 100644 --- a/imperative/python/megengine/distributed/functional.py +++ b/imperative/python/megengine/distributed/functional.py @@ -40,7 +40,7 @@ __all__ = [ ] -@apply.add +@apply.register() def _(op: RemoteSend, *args: Tensor): ret = tensor_apply(op, *args) diff --git a/imperative/python/test/integration/test_dp_correctness.py b/imperative/python/test/integration/test_dp_correctness.py index df94ea4e..8e62105a 100644 --- a/imperative/python/test/integration/test_dp_correctness.py +++ b/imperative/python/test/integration/test_dp_correctness.py @@ -133,7 +133,7 @@ def update_model(model_path): data = Tensor(checkpoint["data"], dtype=np.float32) label = Tensor(checkpoint["label"], dtype=np.int32) - opt.zero_grad() + opt.clear_grad() loss = train(data, label, net=net, opt=opt) opt.step() diff --git a/imperative/python/test/integration/test_optimizer.py b/imperative/python/test/integration/test_optimizer.py index bd7b1798..62cd7b46 100644 --- a/imperative/python/test/integration/test_optimizer.py +++ b/imperative/python/test/integration/test_optimizer.py @@ -73,17 +73,18 @@ def _test_optimizer(opt_str, test_case, check_class, update_lr=False): for symbolic in (False, True): @trace(symbolic=symbolic) - def train_func(data, *, opt=None): - opt.zero_grad() - with opt.record(): + def train_func(data, *, opt=None, gm=None): + opt.clear_grad() + with gm.record(): pred = net(data) loss = pred.sum() - opt.backward(loss) + gm.backward(loss) opt.step() # reset net and opt net = Simple() opt = getattr(optimizer, opt_str)(net.parameters(), **test_case) + gm = ad.GradManager().register(net.parameters()) check_func = check_class(net, **test_case) step = 0 for i in range(iter_num): @@ -96,7 +97,7 @@ def _test_optimizer(opt_str, test_case, check_class, update_lr=False): for param in net.parameters(): ori_params[param] = np.copy(param.numpy()) - train_func(np.random.random(data_shape).astype(np.float32), opt=opt) + train_func(np.random.random(data_shape).astype(np.float32), opt=opt, gm=gm) step += 1 check_func(ori_params, net.parameters(), step) diff --git a/imperative/python/test/integration/test_sgd_momentum.py b/imperative/python/test/integration/test_sgd_momentum.py index f1d75a79..34562e6b 100644 --- a/imperative/python/test/integration/test_sgd_momentum.py +++ b/imperative/python/test/integration/test_sgd_momentum.py @@ -67,23 +67,24 @@ def test_sgd_momentum_trace(): for symbolic in (True, False): @trace(symbolic=symbolic) - def train_func(data, *, model=None, optim=None): - optim.zero_grad() - with optim.record(): + def train_func(data, *, model=None, optim=None, gm=None): + optim.clear_grad() + with gm.record(): loss = net(data) - optim.backward(loss) + gm.backward(loss) optim.step() return loss @trace(symbolic=symbolic) - def eval_func(data, *, model=None, optim=None): + def eval_func(data, *, model=None, optim=None, gm=None): loss = net(data) return loss net = Simple() optim = optimizer.SGD(net.parameters(), lr=1.0, momentum=0.9) + gm = ad.GradManager().register(net.parameters()) data = tensor([2.34]) - train_func(data, model=net, optim=optim) + train_func(data, model=net, optim=optim, gm=gm) np.testing.assert_almost_equal( optim._state[net.a]["momentum_buffer"].numpy(), 2.34 ) @@ -97,7 +98,7 @@ def test_sgd_momentum_trace(): ) # do a step of train - train_func(data, model=net, optim=optim) + train_func(data, model=net, optim=optim, gm=gm) np.testing.assert_almost_equal(loss.numpy(), 2.34 * (1.23 - 2.34), 5) np.testing.assert_almost_equal( optim._state[net.a]["momentum_buffer"].numpy(), 0.9 * 2.34 + 2.34 diff --git a/imperative/python/test/integration/test_trace_dump.py b/imperative/python/test/integration/test_trace_dump.py index 6af911d8..1cd6546a 100644 --- a/imperative/python/test/integration/test_trace_dump.py +++ b/imperative/python/test/integration/test_trace_dump.py @@ -17,6 +17,7 @@ import megengine.functional as F import megengine.module as M import megengine.optimizer as optim from megengine import tensor +from megengine.autodiff import GradManager from megengine.jit import trace @@ -61,17 +62,18 @@ class XORNet(M.Module): def test_xornet_trace_dump(): net = XORNet() opt = optim.SGD(net.parameters(requires_grad=True), lr=0.01, momentum=0.9) + gm = GradManager().register(net.parameters(requires_grad=True)) batch_size = 64 train_dataset = minibatch_generator(batch_size) val_dataset = minibatch_generator(batch_size) @trace def train_fun(data, label): - with opt.record(): + with gm.record(): net.train() pred = net(data) loss = F.cross_entropy_with_softmax(pred, label) - opt.backward(loss) + gm.backward(loss) return pred, loss @trace @@ -95,7 +97,7 @@ def test_xornet_trace_dump(): break data = tensor(minibatch["data"]) label = tensor(minibatch["label"]) - opt.zero_grad() + opt.clear_grad() _, loss = train_fun(data, label) train_loss.append((step, loss.numpy())) if step % 50 == 0: -- GitLab