def __init__(self,
data_covar_module,
num_tasks,
rank=1,
task_covar_prior=None,
bias_only=False,
**kwargs):
"""
"""
super(MyMultitaskKernel, self).__init__(**kwargs)
self.task_covar_module = IndexKernel(num_tasks=num_tasks,
batch_shape=self.batch_shape,
rank=rank,
prior=task_covar_prior)
self.data_covar_module = data_covar_module
self.num_tasks = num_tasks
self.bias_only = bias_only
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,
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)