def test_manual_bounds(self, cuda=False):
device = torch.device("cuda") if cuda else torch.device("cpu")
for dtype in (torch.float, torch.double):
# get a test module
train_x = torch.tensor([[1.0, 2.0, 3.0]], device=device, dtype=dtype)
train_y = torch.tensor([4.0], device=device, dtype=dtype)
likelihood = GaussianLikelihood()
model = ExactGP(train_x, train_y, likelihood)
model.covar_module = RBFKernel(ard_num_dims=3)
model.mean_module = ConstantMean()
model.to(device=device, dtype=dtype)
mll = ExactMarginalLogLikelihood(likelihood, model)
# test the basic case
x, pdict, bounds = module_to_array(
module=mll, bounds={"model.covar_module.raw_lengthscale": (0.1, None)}
)
self.assertTrue(np.array_equal(x, np.zeros(5)))
expected_sizes = {
"likelihood.noise_covar.raw_noise": torch.Size([1]),
"model.covar_module.raw_lengthscale": torch.Size([1, 3]),
"model.mean_module.constant": torch.Size([1]),
}
self.assertEqual(set(pdict.keys()), set(expected_sizes.keys()))
for pname, val in pdict.items():
self.assertEqual(val.dtype, dtype)
self.assertEqual(val.shape, expected_sizes[pname])
self.assertEqual(val.device.type, device.type)
lower_exp = np.full_like(x, 0.1)
for p in ("likelihood.noise_covar.raw_noise", "model.mean_module.constant"):
lower_exp[_get_index(pdict, p)] = -np.inf
self.assertTrue(np.equal(bounds[0], lower_exp).all())
self.assertTrue(np.equal(bounds[1], np.full_like(x, np.inf)).all())
def load_mcmc_samples(self, mcmc_samples: Dict[str, Tensor]) -> None:
r"""Load the MCMC hyperparameter samples into the model.
This method will be called by `fit_fully_bayesian_model_nuts` when the model
has been fitted in order to create a batched SingleTaskGP model.
"""
tkwargs = {"device": self.train_X.device, "dtype": self.train_X.dtype}
num_mcmc_samples = len(mcmc_samples["mean"])
batch_shape = torch.Size([num_mcmc_samples])
self.train_X = self.train_X.unsqueeze(0).expand(
num_mcmc_samples, self.train_X.shape[0], -1
)
self.mean_module = ConstantMean(batch_shape=batch_shape).to(**tkwargs)
self.covar_module = ScaleKernel(
base_kernel=MaternKernel(
ard_num_dims=self.train_X.shape[-1],
batch_shape=batch_shape,
),
batch_shape=batch_shape,
).to(**tkwargs)
if self.train_Yvar is not None:
self.likelihood = FixedNoiseGaussianLikelihood(
noise=self.train_Yvar, batch_shape=batch_shape
).to(**tkwargs)
else:
self.likelihood = GaussianLikelihood(
batch_shape=batch_shape,
noise_constraint=GreaterThan(MIN_INFERRED_NOISE_LEVEL),
).to(**tkwargs)
self.likelihood.noise_covar.noise = (
mcmc_samples["noise"]
.detach()
.clone()
.view(self.likelihood.noise_covar.noise.shape)
.clamp_min(MIN_INFERRED_NOISE_LEVEL)
.to(**tkwargs)
)
self.covar_module.base_kernel.lengthscale = (
mcmc_samples["lengthscale"]
.detach()
.clone()
.view(self.covar_module.base_kernel.lengthscale.shape)
.to(**tkwargs)
)
self.covar_module.outputscale = (
mcmc_samples["outputscale"]
.detach()
.clone()
.view(self.covar_module.outputscale.shape)
.to(**tkwargs)
)
self.mean_module.constant.data = (
mcmc_samples["mean"]
.detach()
.clone()
.view(self.mean_module.constant.shape)
.to(**tkwargs)
)
def test_exclude(self):
for dtype in (torch.float, torch.double):
# get a test module
train_x = torch.tensor([[1.0, 2.0, 3.0]],
device=self.device,
dtype=dtype)
train_y = torch.tensor([4.0], device=self.device, dtype=dtype)
likelihood = GaussianLikelihood()
model = ExactGP(train_x, train_y, likelihood)
model.covar_module = RBFKernel(ard_num_dims=3)
model.mean_module = ConstantMean()
model.to(device=self.device, dtype=dtype)
mll = ExactMarginalLogLikelihood(likelihood, model)
# test the basic case
x, pdict, bounds = module_to_array(
module=mll, exclude={"model.mean_module.constant"})
self.assertTrue(np.array_equal(x, np.zeros(4)))
expected_sizes = {
"likelihood.noise_covar.raw_noise": torch.Size([1]),
"model.covar_module.raw_lengthscale": torch.Size([1, 3]),
}
self.assertEqual(set(pdict.keys()), set(expected_sizes.keys()))
for pname, val in pdict.items():
self.assertEqual(val.dtype, dtype)
self.assertEqual(val.shape, expected_sizes[pname])
self.assertEqual(val.device.type, self.device.type)
self.assertIsNone(bounds)
def __init__(
self,
train_X: Tensor,
train_Y: Tensor,
train_Yvar: Tensor,
outcome_transform: Optional[OutcomeTransform] = None,
) -> None:
r"""A single-task exact GP model using fixed noise levels.
Args:
train_X: A `batch_shape x n x d` tensor of training features.
train_Y: A `batch_shape x n x m` tensor of training observations.
train_Yvar: A `batch_shape x n x m` tensor of observed measurement
noise.
outcome_transform: An outcome transform that is applied to the
training data during instantiation and to the posterior during
inference (that is, the `Posterior` obtained by calling
`.posterior` on the model will be on the original scale).
Example:
>>> train_X = torch.rand(20, 2)
>>> train_Y = torch.sin(train_X).sum(dim=1, keepdim=True)
>>> train_Yvar = torch.full_like(train_Y, 0.2)
>>> model = FixedNoiseGP(train_X, train_Y, train_Yvar)
"""
if outcome_transform is not None:
train_Y, train_Yvar = outcome_transform(train_Y, train_Yvar)
validate_input_scaling(train_X=train_X,
train_Y=train_Y,
train_Yvar=train_Yvar)
self._validate_tensor_args(X=train_X, Y=train_Y, Yvar=train_Yvar)
self._set_dimensions(train_X=train_X, train_Y=train_Y)
train_X, train_Y, train_Yvar = self._transform_tensor_args(
X=train_X, Y=train_Y, Yvar=train_Yvar)
likelihood = FixedNoiseGaussianLikelihood(
noise=train_Yvar, batch_shape=self._aug_batch_shape)
ExactGP.__init__(self,
train_inputs=train_X,
train_targets=train_Y,
likelihood=likelihood)
self.mean_module = ConstantMean(batch_shape=self._aug_batch_shape)
self.covar_module = ScaleKernel(
base_kernel=MaternKernel(
nu=2.5,
ard_num_dims=train_X.shape[-1],
batch_shape=self._aug_batch_shape,
lengthscale_prior=GammaPrior(3.0, 6.0),
),
batch_shape=self._aug_batch_shape,
outputscale_prior=GammaPrior(2.0, 0.15),
)
if outcome_transform is not None:
self.outcome_transform = outcome_transform
self._subset_batch_dict = {
"mean_module.constant": -2,
"covar_module.raw_outputscale": -1,
"covar_module.base_kernel.raw_lengthscale": -3,
}
self.to(train_X)
def __init__(
self,
train_X: Tensor,
train_Y: Tensor,
likelihood: Optional[Likelihood] = None,
covar_module: Optional[Module] = None,
) -> None:
r"""A single-task exact GP model.
Args:
train_X: A `n x d` or `batch_shape x n x d` (batch mode) tensor of training
features.
train_Y: A `n x m` or `batch_shape x n x m` (batch mode) tensor of
training observations.
likelihood: A likelihood. If omitted, use a standard
GaussianLikelihood with inferred noise level.
covar_module: The covariance (kernel) matrix. If omitted, use the
MaternKernel.
Example:
>>> train_X = torch.rand(20, 2)
>>> train_Y = torch.sin(train_X).sum(dim=1, keepdim=True)
>>> model = SingleTaskGP(train_X, train_Y)
"""
validate_input_scaling(train_X=train_X, train_Y=train_Y)
self._validate_tensor_args(X=train_X, Y=train_Y)
self._set_dimensions(train_X=train_X, train_Y=train_Y)
train_X, train_Y, _ = self._transform_tensor_args(X=train_X, Y=train_Y)
if likelihood is None:
noise_prior = GammaPrior(1.1, 0.05)
noise_prior_mode = (noise_prior.concentration -
1) / noise_prior.rate
likelihood = GaussianLikelihood(
noise_prior=noise_prior,
batch_shape=self._aug_batch_shape,
noise_constraint=GreaterThan(
MIN_INFERRED_NOISE_LEVEL,
transform=None,
initial_value=noise_prior_mode,
),
)
else:
self._is_custom_likelihood = True
ExactGP.__init__(self, train_X, train_Y, likelihood)
self.mean_module = ConstantMean(batch_shape=self._aug_batch_shape)
if covar_module is None:
self.covar_module = ScaleKernel(
MaternKernel(
nu=2.5,
ard_num_dims=train_X.shape[-1],
batch_shape=self._aug_batch_shape,
lengthscale_prior=GammaPrior(3.0, 6.0),
),
batch_shape=self._aug_batch_shape,
outputscale_prior=GammaPrior(2.0, 0.15),
)
else:
self.covar_module = covar_module
self.to(train_X)
def __init__(self,
train_X: Tensor,
train_Y: Tensor,
likelihood: Optional[Likelihood] = None) -> None:
r"""A single-task exact GP model.
Args:
train_X: A `n x d` or `batch_shape x n x d` (batch mode) tensor of training
features.
train_Y: A `n x (o)` or `batch_shape x n x (o)` (batch mode) tensor of
training observations.
likelihood: A likelihood. If omitted, use a standard
GaussianLikelihood with inferred noise level.
Example:
>>> train_X = torch.rand(20, 2)
>>> train_Y = torch.sin(train_X[:, 0]) + torch.cos(train_X[:, 1])
>>> model = SingleTaskGP(train_X, train_Y)
"""
ard_num_dims = train_X.shape[-1]
train_X, train_Y, _ = self._set_dimensions(train_X=train_X,
train_Y=train_Y)
train_X, train_Y, _ = multioutput_to_batch_mode_transform(
train_X=train_X, train_Y=train_Y, num_outputs=self._num_outputs)
if likelihood is None:
noise_prior = GammaPrior(1.1, 0.05)
noise_prior_mode = (noise_prior.concentration -
1) / noise_prior.rate
likelihood = GaussianLikelihood(
noise_prior=noise_prior,
batch_shape=self._aug_batch_shape,
noise_constraint=GreaterThan(
MIN_INFERRED_NOISE_LEVEL,
transform=None,
initial_value=noise_prior_mode,
),
)
else:
self._likelihood_state_dict = deepcopy(likelihood.state_dict())
ExactGP.__init__(self, train_X, train_Y, likelihood)
self.mean_module = ConstantMean(batch_shape=self._aug_batch_shape)
self.covar_module = ScaleKernel(
MaternKernel(
nu=2.5,
ard_num_dims=ard_num_dims,
batch_shape=self._aug_batch_shape,
lengthscale_prior=GammaPrior(3.0, 6.0),
),
batch_shape=self._aug_batch_shape,
outputscale_prior=GammaPrior(2.0, 0.15),
)
self.to(train_X)
def __init__(self, train_X: Tensor, train_Y: Tensor, train_Yvar: Tensor) -> None:
r"""A single-task exact GP model using fixed noise levels.
Args:
train_X: A `n x d` or `batch_shape x n x d` (batch mode) tensor of training
features.
train_Y: A `n x (o)` or `batch_shape x n x (o)` (batch mode) tensor of
training observations.
train_Yvar: A `batch_shape x n x (o)` or `batch_shape x n x (o)`
(batch mode) tensor of observed measurement noise.
Example:
>>> train_X = torch.rand(20, 2)
>>> train_Y = torch.sin(train_X[:, 0]]) + torch.cos(train_X[:, 1])
>>> train_Yvar = torch.full_like(train_Y, 0.2)
>>> model = FixedNoiseGP(train_X, train_Y, train_Yvar)
"""
ard_num_dims = train_X.shape[-1]
train_X, train_Y, train_Yvar = self._set_dimensions(
train_X=train_X, train_Y=train_Y, train_Yvar=train_Yvar
)
train_X, train_Y, train_Yvar = multioutput_to_batch_mode_transform(
train_X=train_X,
train_Y=train_Y,
num_outputs=self._num_outputs,
train_Yvar=train_Yvar,
)
likelihood = FixedNoiseGaussianLikelihood(
noise=train_Yvar, batch_shape=self._aug_batch_shape
)
ExactGP.__init__(
self, train_inputs=train_X, train_targets=train_Y, likelihood=likelihood
)
self.mean_module = ConstantMean(batch_shape=self._aug_batch_shape)
self.covar_module = ScaleKernel(
base_kernel=MaternKernel(
nu=2.5,
ard_num_dims=ard_num_dims,
batch_shape=self._aug_batch_shape,
lengthscale_prior=GammaPrior(3.0, 6.0),
),
batch_shape=self._aug_batch_shape,
outputscale_prior=GammaPrior(2.0, 0.15),
)
self.to(train_X)
def test_set_parameters(self):
for dtype in (torch.float, torch.double):
# get a test module
train_x = torch.tensor([[1.0, 2.0, 3.0]],
device=self.device,
dtype=dtype)
train_y = torch.tensor([4.0], device=self.device, dtype=dtype)
likelihood = GaussianLikelihood()
model = ExactGP(train_x, train_y, likelihood)
model.covar_module = RBFKernel(ard_num_dims=3)
model.mean_module = ConstantMean()
model.to(device=self.device, dtype=dtype)
mll = ExactMarginalLogLikelihood(likelihood, model)
# get parameters
x, pdict, bounds = module_to_array(module=mll)
# Set parameters
mll = set_params_with_array(mll,
np.array([1.0, 2.0, 3.0, 4.0,
5.0]), pdict)
z = dict(mll.named_parameters())
self.assertTrue(
torch.equal(
z["likelihood.noise_covar.raw_noise"],
torch.tensor([1.0], device=self.device, dtype=dtype),
))
self.assertTrue(
torch.equal(
z["model.covar_module.raw_lengthscale"],
torch.tensor([[2.0, 3.0, 4.0]],
device=self.device,
dtype=dtype),
))
self.assertTrue(
torch.equal(
z["model.mean_module.constant"],
torch.tensor([5.0], device=self.device, dtype=dtype),
))
# Extract again
x2, pdict2, bounds2 = module_to_array(module=mll)
self.assertTrue(
np.array_equal(x2, np.array([1.0, 2.0, 3.0, 4.0, 5.0])))
def __init__(self, train_X: Tensor, train_Y: Tensor,
train_Yvar: Tensor) -> None:
self._validate_tensor_args(X=train_X, Y=train_Y, Yvar=train_Yvar)
self._set_dimensions(train_X=train_X, train_Y=train_Y)
train_X, train_Y, train_Yvar = self._transform_tensor_args(
X=train_X, Y=train_Y, Yvar=train_Yvar)
likelihood = FixedNoiseGaussianLikelihood(
noise=train_Yvar, batch_shape=self._aug_batch_shape)
ExactGP.__init__(self,
train_inputs=train_X,
train_targets=train_Y,
likelihood=likelihood)
self.mean_module = ConstantMean(batch_shape=self._aug_batch_shape)
self.covar_module = ScaleKernel(
base_kernel=RBFKernel(
ard_num_dims=train_X.shape[-1],
batch_shape=self._aug_batch_shape,
),
batch_shape=self._aug_batch_shape,
)
self.to(train_X)
def __init__(
self,
datapoints: Tensor,
comparisons: Tensor,
covar_module: Optional[Module] = None,
noise_module: Optional[HomoskedasticNoise] = None,
**kwargs,
) -> None:
super().__init__()
r"""A probit-likelihood GP with Laplace approximation model.
A probit-likelihood GP with Laplace approximation model that learns via
pairwise comparison data. By default it uses a scaled-RBF kernel.
Args:
datapoints: A `batch_shape x n x d` tensor of training features.
comparisons: A `batch_shape x m x 2` training comparisons;
comparisons[i] is a noisy indicator suggesting the utility value
of comparisons[i, 0]-th is greater than comparisons[i, 1]-th.
covar_module: Covariance module
noise_module: Noise module
"""
# Compatibility variables with fit_gpytorch_*: Dummy likelihood
# Likelihood is tightly tied with this model and
# it doesn't make much sense to keep it separate
self.likelihood = None
# TODO: remove these variables from `state_dict()` so that when calling
# `load_state_dict()`, only the hyperparameters are copied over
self.register_buffer("datapoints", None)
self.register_buffer("comparisons", None)
self.register_buffer("utility", None)
self.register_buffer("covar_chol", None)
self.register_buffer("likelihood_hess", None)
self.register_buffer("hlcov_eye", None)
self.register_buffer("covar", None)
self.register_buffer("covar_inv", None)
self.train_inputs = []
self.train_targets = None
self.pred_cov_fac_need_update = True
self._input_batch_shape = torch.Size()
self.dim = None
# will be set to match datapoints' dtype and device
# since scipy.optimize.fsolve only works on cpu, it'd be the
# fastest to fit the model on cpu and take samples on gpu to avoid
# overhead of moving data back and forth during fitting time
self.tkwargs = {}
# See set_train_data for additional compatibility variables
self.set_train_data(datapoints, comparisons, update_model=False)
# Set optional parameters
# jitter to add for numerical stability
self._jitter = kwargs.get("jitter", 1e-6)
# Clamping z lim for better numerical stability. See self._calc_z for detail
# norm_cdf(z=3) ~= 0.999, top 0.1% percent
self._zlim = kwargs.get("zlim", 3)
# Stopping creteria in scipy.optimize.fsolve used to find f_map in _update()
# If None, set to 1e-6 by default in _update
self._xtol = kwargs.get("xtol")
# The maximum number of calls to the function in scipy.optimize.fsolve
# If None, set to 100 by default in _update
# If zero, then 100*(N+1) is used by default by fsolve;
self._maxfev = kwargs.get("maxfev")
# Set hyperparameters
# Do not set the batch_shape explicitly so mean_module can operate in both mode
# once fsolve used in _update can run in batch mode, we should explicitly set
# the bacth shape here
self.mean_module = ConstantMean()
# Do not optimize constant mean prior
for param in self.mean_module.parameters():
param.requires_grad = False
# set covariance module
if noise_module is None:
noise_module = HomoskedasticNoise(
noise_prior=SmoothedBoxPrior(-5, 5, 0.5, transform=torch.log),
noise_constraint=GreaterThan(1e-4), # if None, 1e-4 by default
batch_shape=self._input_batch_shape,
)
self.noise_module = noise_module
# set covariance module
if covar_module is None:
ls_prior = GammaPrior(1.2, 0.5)
ls_prior_mode = (ls_prior.concentration - 1) / ls_prior.rate
covar_module = RBFKernel(
batch_shape=self._input_batch_shape,
ard_num_dims=self.dim,
lengthscale_prior=ls_prior,
lengthscale_constraint=Positive(transform=None,
initial_value=ls_prior_mode),
)
self.covar_module = covar_module
self._x0 = None # will store temporary results for warm-starting
if self.datapoints is not None and self.comparisons is not None:
self.to(dtype=self.datapoints.dtype, device=self.datapoints.device)
self._update() # Find f_map for initial parameters
self.to(self.datapoints)
def __init__(
self,
train_X: Tensor,
train_Y: Tensor,
task_feature: int,
covar_module: Optional[Module] = None,
task_covar_prior: Optional[Prior] = None,
output_tasks: Optional[List[int]] = None,
rank: Optional[int] = None,
input_transform: Optional[InputTransform] = None,
outcome_transform: Optional[OutcomeTransform] = None,
) -> None:
r"""Multi-Task GP model using an ICM kernel, inferring observation noise.
Args:
train_X: A `n x (d + 1)` or `b x n x (d + 1)` (batch mode) tensor
of training data. One of the columns should contain the task
features (see `task_feature` argument).
train_Y: A `n x 1` or `b x n x 1` (batch mode) tensor of training
observations.
task_feature: The index of the task feature (`-d <= task_feature <= d`).
output_tasks: A list of task indices for which to compute model
outputs for. If omitted, return outputs for all task indices.
rank: The rank to be used for the index kernel. If omitted, use a
full rank (i.e. number of tasks) kernel.
task_covar_prior : A Prior on the task covariance matrix. Must operate
on p.s.d. matrices. A common prior for this is the `LKJ` prior.
input_transform: An input transform that is applied in the model's
forward pass.
Example:
>>> X1, X2 = torch.rand(10, 2), torch.rand(20, 2)
>>> i1, i2 = torch.zeros(10, 1), torch.ones(20, 1)
>>> train_X = torch.cat([
>>> torch.cat([X1, i1], -1), torch.cat([X2, i2], -1),
>>> ])
>>> train_Y = torch.cat(f1(X1), f2(X2)).unsqueeze(-1)
>>> model = MultiTaskGP(train_X, train_Y, task_feature=-1)
"""
with torch.no_grad():
transformed_X = self.transform_inputs(
X=train_X, input_transform=input_transform)
self._validate_tensor_args(X=transformed_X, Y=train_Y)
all_tasks, task_feature, d = self.get_all_tasks(
transformed_X, task_feature, output_tasks)
if outcome_transform is not None:
train_Y, _ = outcome_transform(train_Y)
# squeeze output dim
train_Y = train_Y.squeeze(-1)
if output_tasks is None:
output_tasks = all_tasks
else:
if set(output_tasks) - set(all_tasks):
raise RuntimeError(
"All output tasks must be present in input data.")
self._output_tasks = output_tasks
self._num_outputs = len(output_tasks)
# TODO (T41270962): Support task-specific noise levels in likelihood
likelihood = GaussianLikelihood(noise_prior=GammaPrior(1.1, 0.05))
# construct indexer to be used in forward
self._task_feature = task_feature
self._base_idxr = torch.arange(d)
self._base_idxr[task_feature:] += 1 # exclude task feature
super().__init__(train_inputs=train_X,
train_targets=train_Y,
likelihood=likelihood)
self.mean_module = ConstantMean()
if covar_module is None:
self.covar_module = ScaleKernel(
base_kernel=MaternKernel(nu=2.5,
ard_num_dims=d,
lengthscale_prior=GammaPrior(
3.0, 6.0)),
outputscale_prior=GammaPrior(2.0, 0.15),
)
else:
self.covar_module = covar_module
num_tasks = len(all_tasks)
self._rank = rank if rank is not None else num_tasks
self.task_covar_module = IndexKernel(num_tasks=num_tasks,
rank=self._rank,
prior=task_covar_prior)
if input_transform is not None:
self.input_transform = input_transform
if outcome_transform is not None:
self.outcome_transform = outcome_transform
self.to(train_X)
def __init__(
self,
train_X: Tensor,
train_Y: Tensor,
likelihood: Optional[MultitaskGaussianLikelihood] = None,
data_covar_module: Optional[Module] = None,
task_covar_prior: Optional[Prior] = None,
rank: Optional[int] = None,
input_transform: Optional[InputTransform] = None,
outcome_transform: Optional[OutcomeTransform] = None,
**kwargs: Any,
) -> None:
r"""Multi-task GP with Kronecker structure, using a simple ICM kernel.
Args:
train_X: A `batch_shape x n x d` tensor of training features.
train_Y: A `batch_shape x n x m` tensor of training observations.
likelihood: A `MultitaskGaussianLikelihood`. If omitted, uses a
`MultitaskGaussianLikelihood` with a `GammaPrior(1.1, 0.05)`
noise prior.
data_covar_module: The module computing the covariance (Kernel) matrix
in data space. If omitted, use a `MaternKernel`.
task_covar_prior : A Prior on the task covariance matrix. Must operate
on p.s.d. matrices. A common prior for this is the `LKJ` prior. If
omitted, uses `LKJCovariancePrior` with `eta` parameter as specified
in the keyword arguments (if not specified, use `eta=1.5`).
rank: The rank of the ICM kernel. If omitted, use a full rank kernel.
kwargs: Additional arguments to override default settings of priors,
including:
- eta: The eta parameter on the default LKJ task_covar_prior.
A value of 1.0 is uninformative, values <1.0 favor stronger
correlations (in magnitude), correlations vanish as eta -> inf.
- sd_prior: A scalar prior over nonnegative numbers, which is used
for the default LKJCovariancePrior task_covar_prior.
- likelihood_rank: The rank of the task covariance matrix to fit.
Defaults to 0 (which corresponds to a diagonal covariance matrix).
Example:
>>> train_X = torch.rand(10, 2)
>>> train_Y = torch.cat([f_1(X), f_2(X)], dim=-1)
>>> model = KroneckerMultiTaskGP(train_X, train_Y)
"""
with torch.no_grad():
transformed_X = self.transform_inputs(
X=train_X, input_transform=input_transform)
if outcome_transform is not None:
train_Y, _ = outcome_transform(train_Y)
self._validate_tensor_args(X=transformed_X, Y=train_Y)
self._num_outputs = train_Y.shape[-1]
batch_shape, ard_num_dims = train_X.shape[:-2], train_X.shape[-1]
num_tasks = train_Y.shape[-1]
if rank is None:
rank = num_tasks
if likelihood is None:
noise_prior = GammaPrior(1.1, 0.05)
noise_prior_mode = (noise_prior.concentration -
1) / noise_prior.rate
likelihood = MultitaskGaussianLikelihood(
num_tasks=num_tasks,
batch_shape=batch_shape,
noise_prior=noise_prior,
noise_constraint=GreaterThan(
MIN_INFERRED_NOISE_LEVEL,
transform=None,
initial_value=noise_prior_mode,
),
rank=kwargs.get("likelihood_rank", 0),
)
if task_covar_prior is None:
task_covar_prior = LKJCovariancePrior(
n=num_tasks,
eta=torch.tensor(kwargs.get("eta", 1.5)).to(train_X),
sd_prior=kwargs.get(
"sd_prior",
SmoothedBoxPrior(math.exp(-6), math.exp(1.25), 0.05),
),
)
super().__init__(train_X, train_Y, likelihood)
self.mean_module = MultitaskMean(
base_means=ConstantMean(batch_shape=batch_shape),
num_tasks=num_tasks)
if data_covar_module is None:
data_covar_module = MaternKernel(
nu=2.5,
ard_num_dims=ard_num_dims,
lengthscale_prior=GammaPrior(3.0, 6.0),
batch_shape=batch_shape,
)
else:
data_covar_module = data_covar_module
self.covar_module = MultitaskKernel(
data_covar_module=data_covar_module,
num_tasks=num_tasks,
rank=rank,
batch_shape=batch_shape,
task_covar_prior=task_covar_prior,
)
if outcome_transform is not None:
self.outcome_transform = outcome_transform
if input_transform is not None:
self.input_transform = input_transform
self.to(train_X)
def __init__(
self,
train_X: Tensor,
train_Y: Tensor,
likelihood: Optional[Likelihood] = None,
covar_module: Optional[Module] = None,
outcome_transform: Optional[OutcomeTransform] = None,
) -> None:
r"""A single-task exact GP model.
Args:
train_X: A `batch_shape x n x d` tensor of training features.
train_Y: A `batch_shape x n x m` tensor of training observations.
likelihood: A likelihood. If omitted, use a standard
GaussianLikelihood with inferred noise level.
covar_module: The module computing the covariance (Kernel) matrix.
If omitted, use a `MaternKernel`.
outcome_transform: An outcome transform that is applied to the
training data during instantiation and to the posterior during
inference (that is, the `Posterior` obtained by calling
`.posterior` on the model will be on the original scale).
Example:
>>> train_X = torch.rand(20, 2)
>>> train_Y = torch.sin(train_X).sum(dim=1, keepdim=True)
>>> model = SingleTaskGP(train_X, train_Y)
"""
if outcome_transform is not None:
train_Y, _ = outcome_transform(train_Y)
validate_input_scaling(train_X=train_X, train_Y=train_Y)
self._validate_tensor_args(X=train_X, Y=train_Y)
self._set_dimensions(train_X=train_X, train_Y=train_Y)
train_X, train_Y, _ = self._transform_tensor_args(X=train_X, Y=train_Y)
if likelihood is None:
noise_prior = GammaPrior(1.1, 0.05)
noise_prior_mode = (noise_prior.concentration -
1) / noise_prior.rate
likelihood = GaussianLikelihood(
noise_prior=noise_prior,
batch_shape=self._aug_batch_shape,
noise_constraint=GreaterThan(
MIN_INFERRED_NOISE_LEVEL,
transform=None,
initial_value=noise_prior_mode,
),
)
else:
self._is_custom_likelihood = True
ExactGP.__init__(self, train_X, train_Y, likelihood)
self.mean_module = ConstantMean(batch_shape=self._aug_batch_shape)
if covar_module is None:
self.covar_module = ScaleKernel(
MaternKernel(
nu=2.5,
ard_num_dims=train_X.shape[-1],
batch_shape=self._aug_batch_shape,
lengthscale_prior=GammaPrior(3.0, 6.0),
),
batch_shape=self._aug_batch_shape,
outputscale_prior=GammaPrior(2.0, 0.15),
)
self._subset_batch_dict = {
"likelihood.noise_covar.raw_noise": -2,
"mean_module.constant": -2,
"covar_module.raw_outputscale": -1,
"covar_module.base_kernel.raw_lengthscale": -3,
}
else:
self.covar_module = covar_module
# TODO: Allow subsetting of other covar modules
if outcome_transform is not None:
self.outcome_transform = outcome_transform
self.to(train_X)
def __init__(
self,
datapoints: Tensor,
comparisons: Tensor,
covar_module: Optional[Module] = None,
input_transform: Optional[InputTransform] = None,
**kwargs,
) -> None:
r"""A probit-likelihood GP with Laplace approximation model that learns via
pairwise comparison data. By default it uses a scaled RBF kernel.
Args:
datapoints: A `batch_shape x n x d` tensor of training features.
comparisons: A `batch_shape x m x 2` training comparisons;
comparisons[i] is a noisy indicator suggesting the utility value
of comparisons[i, 0]-th is greater than comparisons[i, 1]-th.
covar_module: Covariance module.
input_transform: An input transform that is applied in the model's
forward pass.
"""
super().__init__()
if input_transform is not None:
input_transform.to(datapoints)
# input transformation is applied in set_train_data
self.input_transform = input_transform
# Compatibility variables with fit_gpytorch_*: Dummy likelihood
# Likelihood is tightly tied with this model and
# it doesn't make much sense to keep it separate
self.likelihood = None
# TODO: remove these variables from `state_dict()` so that when calling
# `load_state_dict()`, only the hyperparameters are copied over
self.register_buffer("datapoints", None)
self.register_buffer("comparisons", None)
self.register_buffer("D", None)
self.register_buffer("DT", None)
self.register_buffer("utility", None)
self.register_buffer("covar_chol", None)
self.register_buffer("likelihood_hess", None)
self.register_buffer("hlcov_eye", None)
self.register_buffer("covar", None)
self.register_buffer("covar_inv", None)
self.train_inputs = []
self.train_targets = None
self.pred_cov_fac_need_update = True
self.dim = None
# See set_train_data for additional compatibility variables.
# Not that the datapoints here are not transformed even if input_transform
# is not None to avoid double transformation during model fitting.
# self.transform_inputs is called in `forward`
self.set_train_data(datapoints, comparisons, update_model=False)
# Set optional parameters
# jitter to add for numerical stability
self._jitter = kwargs.get("jitter", 1e-6)
# Clamping z lim for better numerical stability. See self._calc_z for detail
# norm_cdf(z=3) ~= 0.999, top 0.1% percent
self._zlim = kwargs.get("zlim", 3)
# Stopping creteria in scipy.optimize.fsolve used to find f_map in _update()
# If None, set to 1e-6 by default in _update
self._xtol = kwargs.get("xtol")
# The maximum number of calls to the function in scipy.optimize.fsolve
# If None, set to 100 by default in _update
# If zero, then 100*(N+1) is used by default by fsolve;
self._maxfev = kwargs.get("maxfev")
# Set hyperparameters
# Do not set the batch_shape explicitly so mean_module can operate in both mode
# once fsolve used in _update can run in batch mode, we should explicitly set
# the bacth shape here
self.mean_module = ConstantMean()
# Do not optimize constant mean prior
for param in self.mean_module.parameters():
param.requires_grad = False
# set covariance module
# the default outputscale here is only a rule of thumb, meant to keep
# estimates away from scale value that would make Phi(f(x)) saturate
# at 0 or 1
if covar_module is None:
ls_prior = GammaPrior(1.2, 0.5)
ls_prior_mode = (ls_prior.concentration - 1) / ls_prior.rate
covar_module = ScaleKernel(
RBFKernel(
batch_shape=self.batch_shape,
ard_num_dims=self.dim,
lengthscale_prior=ls_prior,
lengthscale_constraint=Positive(
transform=None, initial_value=ls_prior_mode),
),
outputscale_prior=SmoothedBoxPrior(a=1, b=4),
)
self.covar_module = covar_module
self._x0 = None # will store temporary results for warm-starting
if self.datapoints is not None and self.comparisons is not None:
self.to(dtype=self.datapoints.dtype, device=self.datapoints.device)
# Find f_map for initial parameters with transformed datapoints
transformed_dp = self.transform_inputs(datapoints)
self._update(transformed_dp)
self.to(self.datapoints)
def __init__(
self,
train_X: Tensor,
train_Y: Tensor,
task_feature: int,
output_tasks: Optional[List[int]] = None,
rank: Optional[int] = None,
) -> None:
r"""Multi-Task GP model using an ICM kernel, inferring observation noise.
Args:
train_X: A `n x (d + 1)` or `b x n x (d + 1)` (batch mode) tensor
of training data. One of the columns should contain the task
features (see `task_feature` argument).
train_Y: A `n` or `b x n` (batch mode) tensor of training
observations.
task_feature: The index of the task feature
(`-d <= task_feature <= d`).
output_tasks: A list of task indices for which to compute model
outputs for. If omitted, return outputs for all task indices.
rank: The rank to be used for the index kernel. If omitted, use a
full rank (i.e. number of tasks) kernel.
Example:
>>> X1, X2 = torch.rand(10, 2), torch.rand(20, 2)
>>> i1, i2 = torch.zeros(10, 1), torch.ones(20, 1)
>>> train_X = torch.stack([
>>> torch.cat([X1, i1], -1), torch.cat([X2, i2], -1),
>>> ])
>>> train_Y = torch.cat(f1(X1), f2(X2))
>>> model = MultiTaskGP(train_X, train_Y, task_feature=-1)
"""
if train_X.ndimension() != 2:
# Currently, batch mode MTGPs are blocked upstream in GPyTorch
raise ValueError(f"Unsupported shape {train_X.shape} for train_X.")
d = train_X.shape[-1] - 1
if not (-d <= task_feature <= d):
raise ValueError(f"Must have that -{d} <= task_feature <= {d}")
all_tasks = train_X[:, task_feature].unique().to(
dtype=torch.long).tolist()
if output_tasks is None:
output_tasks = all_tasks
else:
if any(t not in all_tasks for t in output_tasks):
raise RuntimeError(
"All output tasks must be present in input data.")
self._output_tasks = output_tasks
# TODO (T41270962): Support task-specific noise levels in likelihood
likelihood = GaussianLikelihood(noise_prior=GammaPrior(1.1, 0.05))
# construct indexer to be used in forward
self._task_feature = task_feature
self._base_idxr = torch.arange(d)
self._base_idxr[task_feature:] += 1 # exclude task feature
super().__init__(train_inputs=train_X,
train_targets=train_Y,
likelihood=likelihood)
self.mean_module = ConstantMean()
self.covar_module = ScaleKernel(
base_kernel=MaternKernel(nu=2.5,
ard_num_dims=d,
lengthscale_prior=GammaPrior(3.0, 6.0)),
outputscale_prior=GammaPrior(2.0, 0.15),
)
num_tasks = len(all_tasks)
self._rank = rank if rank is not None else num_tasks
# TODO: Add LKJ prior for the index kernel
self.task_covar_module = IndexKernel(num_tasks=num_tasks,
rank=self._rank)
self.to(train_X)