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

Function _test_optimizer

imperative/python/test/integration/test_optimizer.py:38–113  ·  view source on GitHub ↗
(opt_str, test_case, check_class, update_lr=False)

Source from the content-addressed store, hash-verified

36
37
38def _test_optimizer(opt_str, test_case, check_class, update_lr=False):
39 iter_num = 3
40 net = Simple()
41 opt = getattr(optimizer, opt_str)(net.parameters(), **test_case)
42 check_func = check_class(net, **test_case)
43 gm = ad.GradManager().attach(net.parameters())
44
45 step = 0
46 data_shape = (2, 28)
47
48 for i in range(iter_num):
49 if update_lr and i == 1: # change learning rate
50 for group in opt.param_groups:
51 group["lr"] += 0.01
52 check_func.lr += 0.01
53 data = Tensor(np.random.random(data_shape).astype(np.float32))
54
55 opt.clear_grad()
56 with gm:
57 pred = net(data)
58 loss = pred.sum()
59 gm.backward(loss)
60
61 ori_params = {}
62 ori_grads = {}
63 for param in net.parameters():
64 assert param._tuple_shape is ()
65 ori_params[param] = np.copy(param.numpy())
66 ori_grads[param] = np.copy(param.grad.numpy())
67 opt.step()
68 # check grad not change
69 for param in net.parameters():
70 assert np.equal(
71 ori_grads[param], param.grad.numpy()
72 ), "step should not change param.grad"
73 step += 1
74 check_func(ori_params, net.parameters(), step)
75
76 # static graph
77 for symbolic in (False, True):
78
79 @trace(symbolic=symbolic)
80 def train_func(data, *, opt=None, gm=None):
81 opt.clear_grad()
82 with gm:
83 pred = net(data)
84 loss = pred.sum()
85 gm.backward(loss)
86 opt.step()
87
88 # reset net and opt
89 net = Simple()
90 opt = getattr(optimizer, opt_str)(net.parameters(), **test_case)
91 gm = ad.GradManager().attach(net.parameters())
92 check_func = check_class(net, **test_case)
93 step = 0
94 for i in range(iter_num):
95 if update_lr and i == 1: # change learning rate

Callers 5

test_sgdFunction · 0.85
test_adamFunction · 0.85
test_adagradFunction · 0.85
test_adadeltaFunction · 0.85
test_adamwFunction · 0.85

Calls 14

TensorClass · 0.90
parametersMethod · 0.80
SimpleClass · 0.70
train_funcFunction · 0.70
attachMethod · 0.45
GradManagerMethod · 0.45
astypeMethod · 0.45
clear_gradMethod · 0.45
sumMethod · 0.45
backwardMethod · 0.45
copyMethod · 0.45
numpyMethod · 0.45

Tested by

no test coverage detected