Coregionalization с пакетным обучением в gpflow

Question

Я пытаюсь реализовать модель GP с coregionalization, которая может обучаться, используя партии.

Пример coregionalization в документах (https://gpflow.readthedocs.io/en/master/notebooks/advanced/coregionalisation.html) использует модель VGP (gpflow.models.VGP), и это просто в использовании. Тем не менее, я не могу обучить это с пакетами (это может быть возможно?)

В идеале я хотел бы использовать модель SVGP (gpflow.models.SVGP) - для которой есть хороший пример с пакетным обучением (https://gpflow.readthedocs.io/en/master/notebooks/advanced/gps_for_big_data.html). Модель тренируется (см. Ниже), но я не могу дать ей разумный прогноз. Я думаю, что это предсказание неправильно.

Я запрограммировал расчет «вручную», чтобы получить ожидаемое решение.

Любые идеи очень ценятся!

Большое спасибо,

Уилл.

import gpflow
import tensorflow as tf
import numpy as np
import matplotlib
from tqdm import trange
%matplotlib inline

from copy import copy, deepcopy
from gpflow.ci_utils import ci_niter

matplotlib.rcParams['figure.figsize'] = (12, 6)
plt = matplotlib.pyplot
np.random.seed(123)

Генерирование некоторых данных

# make a dataset with two outputs, correlated, heavy-tail noise. One has more noise than the other.
X1 = np.random.rand(100, 1)  # Observed locations for first output
X2 = np.random.rand(50, 1) * 0.5  # Observed locations for second output

Y1 = np.sin(6*X1) + np.random.randn(*X1.shape) * 0.03
Y2 = np.sin(6*X2 + 0.7) + np.random.randn(*X2.shape) * 0.1

plt.figure(figsize=(8, 4))
plt.plot(X1, Y1, 'x', mew=2)
plt.plot(X2, Y2, 'x', mew=2)

Увеличение данных с многозадачным индексом

# Augment the input with ones or zeros to indicate the required output dimension
X_augmented = np.vstack((np.hstack((X1, np.zeros_like(X1))), np.hstack((X2, np.ones_like(X2)))))

# Augment the Y data
Y_augmented = np.vstack((np.hstack((Y1, np.zeros_like(Y1))), np.hstack((Y2, np.ones_like(Y2)))))

Построение ядра

output_dim = 2  # Number of outputs
rank = 1  # Rank of W

# Base kernel
k = gpflow.kernels.Matern32(active_dims=[0])

# Coregion kernel
coreg = gpflow.kernels.Coregion(output_dim=output_dim, rank=rank, active_dims=[1])

kern = k * coreg

Set до модели

data = (X_augmented, Y_augmented)

M = 50  # Number of inducing locations
N = len(X_augmented)

Z = X_augmented[0::2, :].copy()  # Initialize inducing locations to the first M inputs in the dataset

m = gpflow.models.SVGP(kern, gpflow.likelihoods.Gaussian(), Z, num_data=N, num_latent_gps=1)

gpflow.utilities.set_trainable(m.inducing_variable, False)

настроить партии

minibatch_size = 100

train_dataset = tf.data.Dataset.from_tensor_slices(data) \
    .repeat() \
    .shuffle(N)

train_it = iter(train_dataset.batch(minibatch_size))

Функция обучения помощника

minibatch_size = 50

# We turn off training for inducing point locations
gpflow.utilities.set_trainable(m.inducing_variable, False)

@tf.function(autograph=False)
def optimization_step(optimizer, model: gpflow.models.SVGP, batch):
    with tf.GradientTape(watch_accessed_variables=False) as tape:
        tape.watch(model.trainable_variables)
        objective = - model.elbo(batch)
        grads = tape.gradient(objective, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return objective

def run_adam(model, iterations):
    """
    Utility function running the Adam optimizer

    :param model: GPflow model
    :param interations: number of iterations
    """
    # Create an Adam Optimizer action
    logf = []
    train_it = iter(train_dataset.batch(minibatch_size))
    adam = tf.optimizers.Adam()
    for step in trange(iterations):
        elbo = - optimization_step(adam, model, next(train_it))
        if step % 10 == 0:
            logf.append(elbo.numpy())
    return logf

Модель поезда

maxiter = ci_niter(20000)

logf = run_adam(m, maxiter)
plt.plot(np.arange(maxiter)[::10], logf)
plt.xlabel('iteration')
plt.ylabel('ELBO');

Результат участка

def plot_gp(x, mu, var, color, label):
    plt.plot(x, mu, color=color, lw=2, label=label)
    plt.fill_between(x[:, 0],
                     (mu - 2*np.sqrt(var))[:, 0],
                     (mu + 2*np.sqrt(var))[:, 0],
                     color=color, alpha=0.4)

def plot(m):
    plt.figure(figsize=(8, 4))
    xtest = np.linspace(0, 1, 100)[:, None]
    line, = plt.plot(X1, Y1, 'x', mew=2)
    mu, var = m.predict_f(np.hstack((xtest, np.zeros_like(xtest))))
    plot_gp(xtest, mu, var, line.get_color(), 'Y1')

    line, = plt.plot(X2, Y2, 'x', mew=2)
    mu, var = m.predict_f(np.hstack((xtest, np.ones_like(xtest))))
    plot_gp(xtest, mu, var, line.get_color(), 'Y2')

    plt.legend()

plot(m)

Результат SVGP

Вычисления "вручную"

sig2n = m.likelihood.variance.numpy()

xtest = xtest = np.linspace(0, 1, 100)[:, None]
Xs1 = np.hstack((xtest, np.zeros_like(xtest)))
Xs2 = np.hstack((xtest, np.ones_like(xtest)))
Xs = np.vstack((Xs1,Xs2))

Kmm = kern(Z,Z).numpy()
Kmn = kern(Z,X_augmented).numpy()
Kxm = kern(Xs,Z).numpy()
Kxx = kern(Xs, Xs).numpy()

Sigma = np.linalg.inv(Kmm+1/sig2n*Kmn@Kmn.T)
mu = 1/sig2n * Kmm@Sigma@Kmn@Y_augmented[:,0]
A = Kmm@Sigma@Kmm

B = np.linalg.inv(Kmm)@A@np.linalg.inv(Kmm)

f = Kxm@np.linalg.inv(Kmm)@mu
var = np.diag(Kxx - Kxm@np.linalg.inv(Kmm)@Kxm.T + Kxm@B@Kxm.T)

plt.plot(Xs[:,0], f)

line, = plt.plot(X1, Y1, 'x', mew=2, c='k')
line, = plt.plot(X2, Y2, 'x', mew=2, c='k')
plt.fill_between(xtest[:,0],
                 f[0:100] - 2*np.sqrt(var[0:100]),
                 f[0:100] + 2*np.sqrt(var[0:100]),
                 color='r', alpha=0.4)
plt.fill_between(xtest[:,0],
                 f[100::] - 2*np.sqrt(var[100::]),
                 f[100::] + 2*np.sqrt(var[100::]),
                 color='r', alpha=0.4)

Мой результат