(OptCls, **kwargs)
| 154 | @pytest.mark.skipif(not is_cuda_available(), reason="only support cuda now") |
| 155 | def test_xla_trace_optimizer(): |
| 156 | def tester(OptCls, **kwargs): |
| 157 | def runner(is_trace): |
| 158 | np.random.seed(123) |
| 159 | megengine.random.seed(123) |
| 160 | |
| 161 | model = ConvNet() |
| 162 | model.train() |
| 163 | |
| 164 | image = np.random.randn(3, 8, 3, 32, 32) |
| 165 | label = np.random.randint(0, 10, (3, 8,)) |
| 166 | |
| 167 | gm = autodiff.GradManager().attach(model.parameters()) |
| 168 | optimizer = OptCls(model.parameters(), **kwargs) |
| 169 | |
| 170 | def func(model, optimizer, timage, tlabel): |
| 171 | with gm: |
| 172 | score = model(timage) |
| 173 | loss = F.nn.cross_entropy(score, tlabel) |
| 174 | gm.backward(loss) |
| 175 | optimizer.step().clear_grad() |
| 176 | return loss |
| 177 | |
| 178 | if is_trace: |
| 179 | func = xla_trace(func, without_host=True, capture_as_const=True) |
| 180 | |
| 181 | losses, updated_weights = [], [] |
| 182 | for i in range(6): |
| 183 | timage = megengine.Tensor(image[i % 3]) |
| 184 | tlabel = megengine.Tensor(label[i % 3]) |
| 185 | loss = func(model, optimizer, timage, tlabel) |
| 186 | |
| 187 | losses.append(loss.item()) |
| 188 | updated_weights.append(model.conv1.weight.numpy().reshape(-1)) |
| 189 | |
| 190 | if i == 4: |
| 191 | for pg in optimizer.param_groups: |
| 192 | pg["lr"] = 0.005 |
| 193 | |
| 194 | return np.asarray(losses), np.stack(updated_weights) |
| 195 | |
| 196 | imp_loss, imp_weight = runner(False) |
| 197 | xla_loss, xla_weight = runner(True) |
| 198 | np.testing.assert_allclose(imp_loss, xla_loss, rtol=5e-3) |
| 199 | np.testing.assert_allclose(imp_weight, xla_weight, atol=1e-3) |
| 200 | |
| 201 | tester(SGD, lr=0.01, momentum=0.9, nesterov=True, weight_decay=0.1) |
| 202 | tester(SGD, lr=0.01, momentum=0.9, nesterov=False, weight_decay=0.1) |
no test coverage detected