def _get_model_and_training_data(input_transform=None, **tkwargs):
train_X, train_Y = _get_random_mt_data(**tkwargs)
model = MultiTaskGP(train_X,
train_Y,
task_feature=1,
input_transform=input_transform)
return model.to(**tkwargs), train_X, train_Y
def _get_fixed_prior_model(**tkwargs):
train_X, train_Y = _get_random_mt_data(**tkwargs)
sd_prior = GammaPrior(2.0, 0.15)
sd_prior._event_shape = torch.Size([2])
model = MultiTaskGP(train_X,
train_Y,
task_feature=1,
prior=LKJCovariancePrior(2, 0.6, sd_prior))
return model.to(**tkwargs)
def _get_given_covar_module_model(**tkwargs):
train_X, train_Y = _get_random_mt_data(**tkwargs)
model = MultiTaskGP(
train_X,
train_Y,
task_feature=1,
covar_module=RBFKernel(lengthscale_prior=LogNormalPrior(0.0, 1.0)),
)
return model.to(**tkwargs)
def test_get_gp_samples(self):
# test multi-task model
X = torch.stack([torch.rand(3), torch.tensor([1.0, 0.0, 1.0])], dim=-1)
Y = torch.rand(3, 1)
with self.assertRaises(NotImplementedError):
gp_samples = get_gp_samples(
model=MultiTaskGP(X, Y, task_feature=1),
num_outputs=1,
n_samples=20,
num_rff_features=500,
)
tkwargs = {"device": self.device}
for dtype, m in product((torch.float, torch.double), (1, 2)):
tkwargs["dtype"] = dtype
model, X, Y = _get_model(**tkwargs, multi_output=m == 2)
gp_samples = get_gp_samples(
model=model,
num_outputs=m,
n_samples=20,
num_rff_features=500,
)
self.assertEqual(len(gp_samples), 20)
self.assertIsInstance(gp_samples[0], DeterministicModel)
Y_hat_rff = torch.stack([gp_sample(X) for gp_sample in gp_samples],
dim=0).mean(dim=0)
with torch.no_grad():
Y_hat = model.posterior(X).mean
self.assertTrue(torch.allclose(Y_hat_rff, Y_hat, atol=2e-1))
def _get_model(
X: Tensor,
Y: Tensor,
Yvar: Tensor,
task_feature: Optional[int] = None,
fidelity_features: Optional[List[int]] = None,
**kwargs: Any,
) -> GPyTorchModel:
"""Instantiate a model of type depending on the input data.
Args:
X: A `n x d` tensor of input features.
Y: A `n x m` tensor of input observations.
Yvar: A `n x m` tensor of input variances (NaN if unobserved).
task_feature: The index of the column pertaining to the task feature
(if present).
fidelity_features: List of columns of X that are fidelity parameters.
Returns:
A GPyTorchModel (unfitted).
"""
Yvar = Yvar.clamp_min_(MIN_OBSERVED_NOISE_LEVEL)
is_nan = torch.isnan(Yvar)
any_nan_Yvar = torch.any(is_nan)
all_nan_Yvar = torch.all(is_nan)
if any_nan_Yvar and not all_nan_Yvar:
raise ValueError(
"Mix of known and unknown variances indicates valuation function "
"errors. Variances should all be specified, or none should be."
)
if fidelity_features is None:
fidelity_features = []
if len(fidelity_features) == 0:
# only pass linear_truncated arg if there are fidelities
kwargs = {k: v for k, v in kwargs.items() if k != "linear_truncated"}
if len(fidelity_features) > 0:
if task_feature:
raise NotImplementedError(
"multi-task multi-fidelity models not yet available"
)
# at this point we can assume that there is only a single fidelity parameter
gp = SingleTaskMultiFidelityGP(
train_X=X, train_Y=Y, data_fidelity=fidelity_features[0], **kwargs
)
elif task_feature is None and all_nan_Yvar:
gp = SingleTaskGP(train_X=X, train_Y=Y, **kwargs)
elif task_feature is None:
gp = FixedNoiseGP(train_X=X, train_Y=Y, train_Yvar=Yvar, **kwargs)
elif all_nan_Yvar:
gp = MultiTaskGP(train_X=X, train_Y=Y, task_feature=task_feature, **kwargs)
else:
gp = FixedNoiseMultiTaskGP(
train_X=X,
train_Y=Y.view(-1),
train_Yvar=Yvar.view(-1),
task_feature=task_feature,
**kwargs,
)
return gp
def test_get_gp_samples(self):
# test multi-task model
with torch.random.fork_rng():
torch.manual_seed(0)
X = torch.stack(
[torch.rand(3), torch.tensor([1.0, 0.0, 1.0])], dim=-1)
Y = torch.rand(3, 1)
with self.assertRaises(NotImplementedError):
gp_samples = get_gp_samples(
model=MultiTaskGP(X, Y, task_feature=1),
num_outputs=1,
n_samples=20,
num_rff_features=500,
)
tkwargs = {"device": self.device}
for dtype, m in product((torch.float, torch.double), (1, 2)):
tkwargs["dtype"] = dtype
for mtype in [True, False]:
model, X, Y = _get_model(**tkwargs, multi_output=m == 2)
use_batch_model = mtype and m == 2
with torch.random.fork_rng():
torch.manual_seed(0)
gp_samples = get_gp_samples(
model=batched_to_model_list(model)
if use_batch_model else model,
num_outputs=m,
n_samples=20,
num_rff_features=500,
)
self.assertEqual(len(gp_samples(X)), 20)
self.assertIsInstance(gp_samples, DeterministicModel)
Y_hat_rff = gp_samples(X).mean(dim=0)
with torch.no_grad():
Y_hat = model.posterior(X).mean
self.assertTrue(torch.allclose(Y_hat_rff, Y_hat, atol=2e-1))
# test batched evaluation
Y_batched = gp_samples(
torch.randn(13, 20, 3, X.shape[-1], **tkwargs))
self.assertEqual(Y_batched.shape, torch.Size([13, 20, 3, m]))
# test incorrect batch shape check
with self.assertRaises(ValueError):
gp_samples(torch.randn(13, 23, 3, X.shape[-1], **tkwargs))
def _get_model(
X: Tensor,
Y: Tensor,
Yvar: Tensor,
task_feature: Optional[int],
fidelity_features: Optional[List[int]] = None,
fidelity_model_id: Optional[int] = None,
) -> GPyTorchModel:
"""Instantiate a model of type depending on the input data."""
Yvar = Yvar.clamp_min_(MIN_OBSERVED_NOISE_LEVEL) # pyre-ignore: [16]
is_nan = torch.isnan(Yvar)
any_nan_Yvar = torch.any(is_nan)
all_nan_Yvar = torch.all(is_nan)
if any_nan_Yvar and not all_nan_Yvar:
raise ValueError(
"Mix of known and unknown variances indicates "
"valuation function errors. Variances should all be specified, or "
"none should be.")
if fidelity_features is None:
fidelity_features = []
if fidelity_model_id is None or len(fidelity_features) == 0:
if task_feature is None and all_nan_Yvar:
gp = SingleTaskGP(train_X=X, train_Y=Y)
elif task_feature is None:
gp = FixedNoiseGP(train_X=X, train_Y=Y, train_Yvar=Yvar)
elif all_nan_Yvar:
gp = MultiTaskGP(train_X=X, train_Y=Y, task_feature=task_feature)
else:
gp = FixedNoiseMultiTaskGP(
train_X=X,
train_Y=Y.view(-1),
train_Yvar=Yvar.view(-1),
task_feature=task_feature,
)
else:
gp_model = model_list[fidelity_model_id]
# pyre-ignore [29]
gp = gp_model(train_X=X, train_Y=Y, train_data_fidelity=False)
return gp
def _get_model_single_output(**tkwargs):
train_X, train_Y = _get_random_mt_data(**tkwargs)
model = MultiTaskGP(train_X, train_Y, task_feature=1, output_tasks=[1])
return model.to(**tkwargs)
def test_MultiTaskGP(self):
bounds = torch.tensor([[-1.0, 0.0], [1.0, 1.0]])
for dtype, use_intf in itertools.product(
(torch.float, torch.double), (False, True)
):
tkwargs = {"device": self.device, "dtype": dtype}
intf = (
Normalize(
d=2, bounds=bounds.to(**tkwargs), transform_on_preprocess=True
)
if use_intf
else None
)
model, train_X, _ = _get_model_and_training_data(
input_transform=intf, **tkwargs
)
self.assertIsInstance(model, MultiTaskGP)
self.assertEqual(model.num_outputs, 2)
self.assertIsInstance(model.likelihood, GaussianLikelihood)
self.assertIsInstance(model.mean_module, ConstantMean)
self.assertIsInstance(model.covar_module, ScaleKernel)
matern_kernel = model.covar_module.base_kernel
self.assertIsInstance(matern_kernel, MaternKernel)
self.assertIsInstance(matern_kernel.lengthscale_prior, GammaPrior)
self.assertIsInstance(model.task_covar_module, IndexKernel)
self.assertEqual(model._rank, 2)
self.assertEqual(
model.task_covar_module.covar_factor.shape[-1], model._rank
)
if use_intf:
self.assertIsInstance(model.input_transform, Normalize)
# test model fitting
mll = ExactMarginalLogLikelihood(model.likelihood, model)
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=OptimizationWarning)
mll = fit_gpytorch_model(mll, options={"maxiter": 1}, max_retries=1)
# test posterior
test_x = torch.rand(2, 1, **tkwargs)
posterior_f = model.posterior(test_x)
self.assertIsInstance(posterior_f, GPyTorchPosterior)
self.assertIsInstance(posterior_f.mvn, MultitaskMultivariateNormal)
self.assertEqual(posterior_f.mean.shape, torch.Size([2, 2]))
self.assertEqual(posterior_f.variance.shape, torch.Size([2, 2]))
# check that training data has input transform applied
# check that the train inputs have been transformed and set on the model
if use_intf:
self.assertTrue(
torch.equal(model.train_inputs[0], model.input_transform(train_X))
)
# test that posterior w/ observation noise raises appropriate error
with self.assertRaises(NotImplementedError):
model.posterior(test_x, observation_noise=True)
with self.assertRaises(NotImplementedError):
model.posterior(test_x, observation_noise=torch.rand(2, **tkwargs))
# test posterior w/ single output index
posterior_f = model.posterior(test_x, output_indices=[0])
self.assertIsInstance(posterior_f, GPyTorchPosterior)
self.assertIsInstance(posterior_f.mvn, MultivariateNormal)
self.assertEqual(posterior_f.mean.shape, torch.Size([2, 1]))
self.assertEqual(posterior_f.variance.shape, torch.Size([2, 1]))
# test posterior w/ bad output index
with self.assertRaises(ValueError):
model.posterior(test_x, output_indices=[2])
# test posterior (batch eval)
test_x = torch.rand(3, 2, 1, **tkwargs)
posterior_f = model.posterior(test_x)
self.assertIsInstance(posterior_f, GPyTorchPosterior)
self.assertIsInstance(posterior_f.mvn, MultitaskMultivariateNormal)
# test that unsupported batch shape MTGPs throw correct error
with self.assertRaises(ValueError):
MultiTaskGP(torch.rand(2, 2, 2), torch.rand(2, 2, 1), 0)
# test that bad feature index throws correct error
train_X, train_Y = _get_random_mt_data(**tkwargs)
with self.assertRaises(ValueError):
MultiTaskGP(train_X, train_Y, 2)
# test that bad output task throws correct error
with self.assertRaises(RuntimeError):
MultiTaskGP(train_X, train_Y, 0, output_tasks=[2])
# test error if outcome_transform attribute is present
model.outcome_transform = None
with self.assertRaises(NotImplementedError):
model.posterior(test_x)
def test_MultiTaskGP(self):
for dtype in (torch.float, torch.double):
tkwargs = {"device": self.device, "dtype": dtype}
model = _get_model(**tkwargs)
self.assertIsInstance(model, MultiTaskGP)
self.assertEqual(model.num_outputs, 2)
self.assertIsInstance(model.likelihood, GaussianLikelihood)
self.assertIsInstance(model.mean_module, ConstantMean)
self.assertIsInstance(model.covar_module, ScaleKernel)
matern_kernel = model.covar_module.base_kernel
self.assertIsInstance(matern_kernel, MaternKernel)
self.assertIsInstance(matern_kernel.lengthscale_prior, GammaPrior)
self.assertIsInstance(model.task_covar_module, IndexKernel)
self.assertEqual(model._rank, 2)
self.assertEqual(model.task_covar_module.covar_factor.shape[-1],
model._rank)
# test model fitting
mll = ExactMarginalLogLikelihood(model.likelihood, model)
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=OptimizationWarning)
mll = fit_gpytorch_model(mll,
options={"maxiter": 1},
max_retries=1)
# test posterior
test_x = torch.rand(2, 1, **tkwargs)
posterior_f = model.posterior(test_x)
self.assertIsInstance(posterior_f, GPyTorchPosterior)
self.assertIsInstance(posterior_f.mvn, MultitaskMultivariateNormal)
self.assertEqual(posterior_f.mean.shape, torch.Size([2, 2]))
self.assertEqual(posterior_f.variance.shape, torch.Size([2, 2]))
# test that posterior w/ observation noise raises appropriate error
with self.assertRaises(NotImplementedError):
model.posterior(test_x, observation_noise=True)
with self.assertRaises(NotImplementedError):
model.posterior(test_x,
observation_noise=torch.rand(2, **tkwargs))
# test posterior w/ single output index
posterior_f = model.posterior(test_x, output_indices=[0])
self.assertIsInstance(posterior_f, GPyTorchPosterior)
self.assertIsInstance(posterior_f.mvn, MultivariateNormal)
self.assertEqual(posterior_f.mean.shape, torch.Size([2, 1]))
self.assertEqual(posterior_f.variance.shape, torch.Size([2, 1]))
# test posterior w/ bad output index
with self.assertRaises(ValueError):
model.posterior(test_x, output_indices=[2])
# test posterior (batch eval)
test_x = torch.rand(3, 2, 1, **tkwargs)
posterior_f = model.posterior(test_x)
self.assertIsInstance(posterior_f, GPyTorchPosterior)
self.assertIsInstance(posterior_f.mvn, MultitaskMultivariateNormal)
# test that unsupported batch shape MTGPs throw correct error
with self.assertRaises(ValueError):
MultiTaskGP(torch.rand(2, 2, 2), torch.rand(2, 1), 0)
# test that bad feature index throws correct error
train_X, train_Y = _get_random_mt_data(**tkwargs)
with self.assertRaises(ValueError):
MultiTaskGP(train_X, train_Y, 2)
# test that bad output task throws correct error
with self.assertRaises(RuntimeError):
MultiTaskGP(train_X, train_Y, 0, output_tasks=[2])
# test error if outcome_transform attribute is present
model.outcome_transform = None
with self.assertRaises(NotImplementedError):
model.posterior(test_x)
def _get_model(**tkwargs):
train_X, train_Y = _get_random_mt_data(**tkwargs)
model = MultiTaskGP(train_X, train_Y, task_feature=1)
return model.to(**tkwargs)
def _get_model(
X: Tensor,
Y: Tensor,
Yvar: Tensor,
task_feature: Optional[int] = None,
fidelity_features: Optional[List[int]] = None,
use_input_warping: bool = False,
**kwargs: Any,
) -> GPyTorchModel:
"""Instantiate a model of type depending on the input data.
Args:
X: A `n x d` tensor of input features.
Y: A `n x m` tensor of input observations.
Yvar: A `n x m` tensor of input variances (NaN if unobserved).
task_feature: The index of the column pertaining to the task feature
(if present).
fidelity_features: List of columns of X that are fidelity parameters.
Returns:
A GPyTorchModel (unfitted).
"""
Yvar = Yvar.clamp_min(MIN_OBSERVED_NOISE_LEVEL) # pyre-ignore[16]
is_nan = torch.isnan(Yvar)
any_nan_Yvar = torch.any(is_nan)
all_nan_Yvar = torch.all(is_nan)
if any_nan_Yvar and not all_nan_Yvar:
if task_feature:
# TODO (jej): Replace with inferred noise before making perf judgements.
Yvar[Yvar != Yvar] = MIN_OBSERVED_NOISE_LEVEL
else:
raise ValueError(
"Mix of known and unknown variances indicates valuation function "
"errors. Variances should all be specified, or none should be."
)
if use_input_warping:
warp_tf = get_warping_transform(
d=X.shape[-1],
task_feature=task_feature,
batch_shape=X.shape[:-2], # pyre-ignore [6]
)
else:
warp_tf = None
if fidelity_features is None:
fidelity_features = []
if len(fidelity_features) == 0:
# only pass linear_truncated arg if there are fidelities
kwargs = {k: v for k, v in kwargs.items() if k != "linear_truncated"}
if len(fidelity_features) > 0:
if task_feature:
raise NotImplementedError( # pragma: no cover
"multi-task multi-fidelity models not yet available"
)
# at this point we can assume that there is only a single fidelity parameter
gp = SingleTaskMultiFidelityGP(
train_X=X,
train_Y=Y,
data_fidelity=fidelity_features[0],
input_transform=warp_tf,
**kwargs,
)
elif task_feature is None and all_nan_Yvar:
gp = SingleTaskGP(train_X=X, train_Y=Y, input_transform=warp_tf, **kwargs)
elif task_feature is None:
gp = FixedNoiseGP(
train_X=X, train_Y=Y, train_Yvar=Yvar, input_transform=warp_tf, **kwargs
)
else:
# instantiate multitask GP
all_tasks, _, _ = MultiTaskGP.get_all_tasks(X, task_feature)
num_tasks = len(all_tasks)
prior_dict = kwargs.get("prior")
prior = None
if prior_dict is not None:
prior_type = prior_dict.get("type", None)
if issubclass(prior_type, LKJCovariancePrior):
sd_prior = prior_dict.get("sd_prior", GammaPrior(1.0, 0.15))
sd_prior._event_shape = torch.Size([num_tasks])
eta = prior_dict.get("eta", 0.5)
if not isinstance(eta, float) and not isinstance(eta, int):
raise ValueError(f"eta must be a real number, your eta was {eta}")
prior = LKJCovariancePrior(num_tasks, eta, sd_prior)
else:
raise NotImplementedError(
"Currently only LKJ prior is supported,"
f"your prior type was {prior_type}."
)
if all_nan_Yvar:
gp = MultiTaskGP(
train_X=X,
train_Y=Y,
task_feature=task_feature,
rank=kwargs.get("rank"),
task_covar_prior=prior,
input_transform=warp_tf,
)
else:
gp = FixedNoiseMultiTaskGP(
train_X=X,
train_Y=Y,
train_Yvar=Yvar,
task_feature=task_feature,
rank=kwargs.get("rank"),
task_covar_prior=prior,
input_transform=warp_tf,
)
return gp
def test_MultiTaskGP(self, cuda=False):
for double in (False, True):
tkwargs = {
"device":
torch.device("cuda") if cuda else torch.device("cpu"),
"dtype": torch.double if double else torch.float,
}
model = _get_model(**tkwargs)
self.assertIsInstance(model, MultiTaskGP)
self.assertIsInstance(model.likelihood, GaussianLikelihood)
self.assertIsInstance(model.mean_module, ConstantMean)
self.assertIsInstance(model.covar_module, ScaleKernel)
matern_kernel = model.covar_module.base_kernel
self.assertIsInstance(matern_kernel, MaternKernel)
self.assertIsInstance(matern_kernel.lengthscale_prior, GammaPrior)
self.assertIsInstance(model.task_covar_module, IndexKernel)
self.assertEqual(model._rank, 2)
self.assertEqual(model.task_covar_module.covar_factor.shape[-1],
model._rank)
# test model fitting
mll = ExactMarginalLogLikelihood(model.likelihood, model)
mll = fit_gpytorch_model(mll, options={"maxiter": 1})
# test posterior
test_x = torch.rand(2, 1, **tkwargs)
posterior_f = model.posterior(test_x)
self.assertIsInstance(posterior_f, GPyTorchPosterior)
self.assertIsInstance(posterior_f.mvn, MultitaskMultivariateNormal)
self.assertEqual(posterior_f.mean.shape, torch.Size([2, 2]))
self.assertEqual(posterior_f.variance.shape, torch.Size([2, 2]))
# test posterior w/ observation noise
posterior_o = model.posterior(test_x, observation_noise=True)
self.assertIsInstance(posterior_o, GPyTorchPosterior)
self.assertIsInstance(posterior_f.mvn, MultitaskMultivariateNormal)
self.assertEqual(posterior_f.mean.shape, torch.Size([2, 2]))
self.assertEqual(posterior_f.variance.shape, torch.Size([2, 2]))
# test posterior w/ single output index
posterior_f = model.posterior(test_x, output_indices=[0])
self.assertIsInstance(posterior_f, GPyTorchPosterior)
self.assertIsInstance(posterior_f.mvn, MultivariateNormal)
self.assertEqual(posterior_f.mean.shape, torch.Size([2, 1]))
self.assertEqual(posterior_f.variance.shape, torch.Size([2, 1]))
# test posterior w/ bad output index
with self.assertRaises(ValueError):
model.posterior(test_x, output_indices=[2])
# test posterior (batch eval)
test_x = torch.rand(3, 2, 1, **tkwargs)
posterior_f = model.posterior(test_x)
self.assertIsInstance(posterior_f, GPyTorchPosterior)
self.assertIsInstance(posterior_f.mvn, MultitaskMultivariateNormal)
# test that unsupported batch shape MTGPs throw correct error
with self.assertRaises(ValueError):
MultiTaskGP(torch.rand(2, 2, 2), torch.rand(2, 1), 0)
# test that bad feature index throws correct error
train_X, train_Y = _get_random_mt_data(**tkwargs)
with self.assertRaises(ValueError):
MultiTaskGP(train_X, train_Y, 2)
# test that bad output task throws correct error
with self.assertRaises(RuntimeError):
MultiTaskGP(train_X, train_Y, 0, output_tasks=[2])
def _get_model_single_output(**tkwargs):
train_X, train_Y = _get_random_mt_data(**tkwargs)
model = MultiTaskGP(train_X, train_Y, task_feature=1, output_tasks=[1])
return model.to(**tkwargs)
def _get_model(**tkwargs):
train_X, train_Y = _get_random_mt_data(**tkwargs)
model = MultiTaskGP(train_X, train_Y, task_feature=1)
return model.to(**tkwargs)
