def update_model(model_path):
"""
Update the dumped model with test cases for new reference values.
The model with pre-trained weights is trained for one iter with the test data attached.
The loss and updated net state dict is dumped.
.. code-block:: python
from test_correctness import update_model
update_model('mnist_model_with_test.mge') # for gpu
update_model('mnist_model_with_test_cpu.mge') # for cpu
"""
net = MnistNet(has_bn=True)
checkpoint = mge.load(model_path)
net.load_state_dict(checkpoint["net_init"])
lr = checkpoint["sgd_lr"]
opt = SGD(net.parameters(), lr=lr)
gm = ad.GradManager().attach(net.parameters())
data = Tensor(checkpoint["data"], dtype=np.float32)
label = Tensor(checkpoint["label"], dtype=np.int32)
opt.clear_grad()
loss = train(data, label, net, opt, gm)
opt.step()
xpu_name = get_xpu_name()
checkpoint.update(
{"net_updated": net.state_dict(), "loss": loss.numpy(), "xpu": xpu_name}
)
mge.save(checkpoint, model_path)
def run_train(
model_path,
use_jit,
use_symbolic,
sublinear_memory_config=None,
max_err=None,
use_adaptive_pooling=False,
):
"""
Load the model with test cases and run the training for one iter.
The loss and updated weights are compared with reference value to verify the correctness.
Dump a new file with updated result by calling update_model
if you think the test fails due to numerical rounding errors instead of bugs.
Please think twice before you do so.
"""
net = MnistNet(has_bn=True, use_adaptive_pooling=use_adaptive_pooling)
checkpoint = mge.load(model_path)
net.load_state_dict(checkpoint["net_init"])
lr = checkpoint["sgd_lr"]
opt = SGD(net.parameters(), lr=lr)
gm = ad.GradManager().attach(net.parameters())
data = Tensor(checkpoint["data"], dtype=np.float32)
label = Tensor(checkpoint["label"], dtype=np.int32)
if max_err is None:
max_err = 1e-5
train_func = train
if use_jit:
train_func = jit.trace(
train_func,
symbolic=use_symbolic,
sublinear_memory_config=sublinear_memory_config,
)
opt.clear_grad()
loss = train_func(data, label, net, opt, gm)
opt.step()
np.testing.assert_allclose(loss.numpy(), checkpoint["loss"], atol=max_err)
for param, param_ref in zip(
net.state_dict().items(), checkpoint["net_updated"].items()
):
assert param[0] == param_ref[0]
if "bn" in param[0]:
ref = param_ref[1].reshape(param[1].shape)
np.testing.assert_allclose(param[1], ref, atol=max_err)
else:
np.testing.assert_allclose(param[1], param_ref[1], atol=max_err)
def test_training_converge_with_swap_and_drop():
_set_swap_flag(True)
_set_drop_flag(True)
old_buffer_length = get_option("buffer_length")
set_option("buffer_length", 0)
net = XORNet()
opt = SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)
gm = ad.GradManager().attach(net.parameters())
def train(data, label):
with gm:
pred = net(data)
loss = F.nn.cross_entropy(pred, label)
gm.backward(loss)
return loss
def infer(data):
return net(data)
train_dataset = minibatch_generator()
losses = []
for data, label in itertools.islice(train_dataset, 2000):
data = Tensor(data, dtype=np.float32)
label = Tensor(label, dtype=np.int32)
opt.clear_grad()
loss = train(data, label)
opt.step()
losses.append(loss.numpy())
assert np.mean(
losses[-100:]) < 0.1, "Final training Loss must be low enough"
ngrid = 10
x = np.linspace(-1.0, 1.0, ngrid)
xx, yy = np.meshgrid(x, x)
xx = xx.reshape((ngrid * ngrid, 1))
yy = yy.reshape((ngrid * ngrid, 1))
data = np.concatenate((xx, yy), axis=1).astype(np.float32)
pred = infer(Tensor(data)).numpy()
precision = calculate_precision(data, pred)
assert precision == 1.0, "Test precision must be high enough, get {}".format(
precision)
_set_swap_flag(False)
_set_drop_flag(False)
set_option("buffer_length", old_buffer_length)
def worker():
net = Simple()
opt = SGD(net.parameters(), lr=0.1)
gm = ad.GradManager().attach(
net.parameters(),
callbacks=[dist.make_allreduce_cb("MEAN", dist.WORLD)])
opt.clear_grad()
with gm:
x = tensor(data)
loss = net(x)
loss = loss.sum()
gm.backward(loss)
for p in net.params:
np.testing.assert_equal(p.grad.numpy(), 1)
def test_training_converge(test_traced_module):
net = XORNet()
if test_traced_module:
inp = Tensor(np.random.random((14, 2)))
net = trace_module(net, inp)
opt = SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)
gm = ad.GradManager().attach(net.parameters())
@trace(symbolic=False)
def train(data, label):
with gm:
pred = net(data)
loss = F.nn.cross_entropy(pred, label)
gm.backward(loss)
optim.clip_grad_norm(net.parameters(), max_norm=0.2, ord=2.0)
return loss
def infer(data):
return net(data)
train_dataset = minibatch_generator()
losses = []
for data, label in itertools.islice(train_dataset, 2000):
data = Tensor(data, dtype=np.float32)
label = Tensor(label, dtype=np.int32)
opt.clear_grad()
loss = train(data, label)
optim.clip_grad_value(net.parameters(), lower=-0.1, upper=0.1)
opt.step()
losses.append(loss.numpy())
assert (np.mean(losses[-100:]) <
0.1), "Final training Loss must be low enough, get {}".format(
np.mean(losses[-100:]))
ngrid = 10
x = np.linspace(-1.0, 1.0, ngrid)
xx, yy = np.meshgrid(x, x)
xx = xx.reshape((ngrid * ngrid, 1))
yy = yy.reshape((ngrid * ngrid, 1))
data = mge.tensor(np.concatenate((xx, yy), axis=1).astype(np.float32))
pred = infer(data)
precision = calculate_precision(data.numpy(), pred.numpy())
assert precision == 1.0, "Test precision must be high enough, get {}".format(
precision)
