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 inste
(
model_path, use_jit, use_symbolic, sublinear_memory_config=None, max_err=None,
)
| 130 | |
| 131 | |
| 132 | def run_test( |
| 133 | model_path, use_jit, use_symbolic, sublinear_memory_config=None, max_err=None, |
| 134 | ): |
| 135 | |
| 136 | """ |
| 137 | Load the model with test cases and run the training for one iter. |
| 138 | The loss and updated weights are compared with reference value to verify the correctness. |
| 139 | |
| 140 | Dump a new file with updated result by calling update_model |
| 141 | if you think the test fails due to numerical rounding errors instead of bugs. |
| 142 | Please think twice before you do so. |
| 143 | |
| 144 | """ |
| 145 | checkpoint = mge.load(model_path) |
| 146 | data = checkpoint["data"] |
| 147 | label = checkpoint["label"] |
| 148 | |
| 149 | @dist.launcher |
| 150 | def worker(max_err): |
| 151 | net = MnistNet(has_bn=True) |
| 152 | net.load_state_dict(checkpoint["net_init"]) |
| 153 | lr = checkpoint["sgd_lr"] |
| 154 | opt = SGD(net.parameters(), lr=lr) |
| 155 | |
| 156 | gm = ad.GradManager().attach( |
| 157 | net.parameters(), callbacks=[dist.make_allreduce_cb("MEAN", dist.WORLD)] |
| 158 | ) |
| 159 | |
| 160 | # use same data and label for all gpu's |
| 161 | # such that the result does not depend on number of gpu |
| 162 | data_train = Tensor(data) |
| 163 | label_train = Tensor(label) |
| 164 | |
| 165 | loss = train(data_train, label_train, net, opt, gm) |
| 166 | |
| 167 | np.testing.assert_allclose(loss.numpy(), checkpoint["loss"], atol=max_err) |
| 168 | |
| 169 | if dist.get_rank(): |
| 170 | return |
| 171 | for param, param_ref in zip( |
| 172 | net.state_dict().items(), checkpoint["net_updated"].items() |
| 173 | ): |
| 174 | assert param[0] == param_ref[0] |
| 175 | if "bn" in param[0]: |
| 176 | ref = param_ref[1].reshape(param[1].shape) |
| 177 | np.testing.assert_allclose(param[1], ref, atol=max_err) |
| 178 | else: |
| 179 | np.testing.assert_allclose(param[1], param_ref[1], atol=max_err) |
| 180 | |
| 181 | worker(max_err) |
| 182 | |
| 183 | |
| 184 | @pytest.mark.require_ngpu(2) |
no test coverage detected