diff --git a/xopt/tests/generators/bayesian/test_model_constructor.py b/xopt/tests/generators/bayesian/test_model_constructor.py index b3f043d7..e4caed16 100644 --- a/xopt/tests/generators/bayesian/test_model_constructor.py +++ b/xopt/tests/generators/bayesian/test_model_constructor.py @@ -995,7 +995,7 @@ def test_train_model_batch_compare(self): for i in range(test_vocs.n_outputs): assert torch.allclose(list_ls[i], batch_ls[i, ...], rtol=0, atol=5e-3) - def test_train_model_batch_compare_adam(self): + def test_train_model_batch_compare_training(self): test_vocs = deepcopy(TEST_VOCS) test_data = deepcopy(TEST_DATA) @@ -1020,17 +1020,48 @@ def test_train_model_batch_compare_adam(self): batch_ls = model_single.covar_module.raw_lengthscale for i in range(test_vocs.n_outputs): assert torch.allclose(list_ls[i], batch_ls[i, ...], rtol=0, atol=1e-3) - print([(p, p.shape) for p in model_list.parameters()]) - print("------------------------------") - print([(p, p.shape) for p in model_single.parameters()]) - optimizer_single = torch.optim.Adam(model_single.parameters(), lr=0.1) mll = ExactMarginalLogLikelihood(model_single.likelihood, model_single) - optimizer_list = [] - mll_list = [] - for ml in model_list.models: - optimizer_list.append(torch.optim.Adam(ml.parameters(), lr=0.1)) - mll_list.append(ExactMarginalLogLikelihood(ml.likelihood, ml)) + mll_list = [ + ExactMarginalLogLikelihood(ml.likelihood, ml) for ml in model_list.models + ] + + # For independent outputs, the batched MLL is mathematically the sum of + # the per-model MLLs. Compare losses and parameter gradients before + # optimizer state can amplify numerical differences. + model_single.zero_grad() + loss = -mll( + model_single(model_single.train_inputs[0]), model_single.train_targets + ).sum() + loss.backward() + single_loss = float(loss.item()) + + total_list_loss = 0.0 + for j, ml in enumerate(model_list.models): + ml.zero_grad() + list_loss = -mll_list[j](ml(ml.train_inputs[0]), ml.train_targets) + list_loss.backward() + total_list_loss += list_loss.item() + assert np.isclose(total_list_loss, single_loss, rtol=0.0, atol=1e-9) + + batch_params = dict(model_single.named_parameters()) + for j, ml in enumerate(model_list.models): + for name, p_list in ml.named_parameters(): + grad_batch = batch_params[name].grad + if grad_batch.shape != p_list.grad.shape: + assert grad_batch.shape[0] == test_vocs.n_outputs, name + grad_batch = grad_batch[j] + assert torch.allclose(grad_batch, p_list.grad, rtol=0, atol=1e-9), name + + # Use SGD for the trajectory check so batched-vs-per-model rounding + # differences remain proportional to the gradient differences. Adam can + # amplify near-zero differences through its adaptive denominator. + # Keep the step count and learning rate small to avoid ill-conditioned + # covariance matrices. + optimizer_single = torch.optim.SGD(model_single.parameters(), lr=0.1) + optimizer_list = [ + torch.optim.SGD(ml.parameters(), lr=0.1) for ml in model_list.models + ] for i in range(10): optimizer_single.zero_grad() @@ -1039,7 +1070,6 @@ def test_train_model_batch_compare_adam(self): loss.backward() optimizer_single.step() single_loss = float(loss.item()) - print(f"Single: {loss.item()}") total_list_loss = 0.0 for j, ml in enumerate(model_list.models): optimizer_list[j].zero_grad() @@ -1047,17 +1077,15 @@ def test_train_model_batch_compare_adam(self): loss = -mll_list[j](output, ml.train_targets) loss.backward() optimizer_list[j].step() - print(f"List {j}: {loss.item()}") total_list_loss += loss.item() - print(f"List total: {total_list_loss}") - assert np.isclose(total_list_loss, single_loss, rtol=0.0, atol=1e-5) + assert np.isclose(total_list_loss, single_loss, rtol=0.0, atol=1e-8) list_ls = [ model_list.models[i].covar_module.raw_lengthscale for i in range(test_vocs.n_outputs) ] batch_ls = model_single.covar_module.raw_lengthscale for i in range(test_vocs.n_outputs): - assert torch.allclose(list_ls[i], batch_ls[i, ...], rtol=0, atol=1e-3) + assert torch.allclose(list_ls[i], batch_ls[i, ...], rtol=0, atol=1e-8) @pytest.mark.parametrize("use_cuda", cuda_combinations) def test_batched_basic(self, use_cuda):