LMC task_indices with pyro

[michaelcao28]

Hi all,

I am trying to combine the use of task_indices for LMC kernel as shown here, but with pyro integration. I need to integrate pyro with gpytorch for my project. What I need to do next is to allow for each data output to have a different set of inputs.

In the tutorial, each epochs_iter looks something like this:

task_indices = torch.LongTensor([indices])
mll = gpytorch.mlls.VariationalELBO(likelihood, model, num_data=train_y.size(0))
optimizer = torch.optim.Adam([
    {'params': model.parameters()},
    {'params': likelihood.parameters()},
], lr=0.1)
...
for i in epochs_iter:
    optimizer.zero_grad()
    output = model(x, task_indices=task_indices)
    loss = -mll(output, train_y)
    loss.backward()
    optimizer.step()

With pyro integration, each epochs_iter looks like this:

elbo = pyro.infer.Trace_ELBO(vectorize_particles=True, retain_graph=True)
optimizer = pyro.optim.Adam({"lr": 0.1})
svi = pyro.infer.SVI(model.model, model.guide, optimizer, elbo)
...
for i in iterator:
        model.zero_grad()
        loss = svi.step(train_x, train_y)

Any ideas how one can implement this?

Many thanks!

[gpleiss]

You should just be able to pass in task_indices=task_indices into the SVI step (i.e. svi.step(train_x, train_y, task_indices=task_indices)). Both model.model and model.guide accept and pass in variable-length argument and warts (see here in code: https://github.com/cornellius-gp/gpytorch/blob/master/gpytorch/models/pyro/pyro_gp.py)