Mixing TensorFlow models with GPflow#
This notebook explores the combination of Keras TensorFlow neural networks with GPflow models.
[1]:
from typing import Dict, Optional, Tuple
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
from scipy.cluster.vq import kmeans2
import gpflow
from gpflow.ci_utils import reduce_in_tests
from gpflow.utilities import to_default_float
iterations = reduce_in_tests(100)
2023-05-03 23:19:19.356754: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations: AVX2 AVX512F FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-05-03 23:19:19.480593: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2023-05-03 23:19:19.480614: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.
2023-05-03 23:19:19.509272: E tensorflow/stream_executor/cuda/cuda_blas.cc:2981] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2023-05-03 23:19:20.119363: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2023-05-03 23:19:20.119423: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2023-05-03 23:19:20.119431: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
Convolutional network inside a GPflow model#
[2]:
original_dataset, info = tfds.load(
name="mnist", split=tfds.Split.TRAIN, with_info=True
)
total_num_data = info.splits["train"].num_examples
image_shape = info.features["image"].shape
image_size = tf.reduce_prod(image_shape)
batch_size = 32
def map_fn(input_slice: Dict[str, tf.Tensor]):
updated = input_slice
image = to_default_float(updated["image"]) / 255.0
label = to_default_float(updated["label"])
return tf.reshape(image, [-1, image_size]), label
autotune = tf.data.experimental.AUTOTUNE
dataset = (
original_dataset.shuffle(1024)
.batch(batch_size, drop_remainder=True)
.map(map_fn, num_parallel_calls=autotune)
.prefetch(autotune)
.repeat()
)
2023-05-03 23:19:23.192026: W tensorflow/core/platform/cloud/google_auth_provider.cc:184] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with "NOT_FOUND: Could not locate the credentials file.". Retrieving token from GCE failed with "FAILED_PRECONDITION: Error executing an HTTP request: libcurl code 6 meaning 'Couldn't resolve host name', error details: Could not resolve host: metadata".
Downloading and preparing dataset 11.06 MiB (download: 11.06 MiB, generated: 21.00 MiB, total: 32.06 MiB) to /home/circleci/tensorflow_datasets/mnist/3.0.1...
Dataset mnist downloaded and prepared to /home/circleci/tensorflow_datasets/mnist/3.0.1. Subsequent calls will reuse this data.
2023-05-03 23:19:23.882574: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory
2023-05-03 23:19:23.882595: W tensorflow/stream_executor/cuda/cuda_driver.cc:263] failed call to cuInit: UNKNOWN ERROR (303)
2023-05-03 23:19:23.882614: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:156] kernel driver does not appear to be running on this host (aeab2aa79215): /proc/driver/nvidia/version does not exist
2023-05-03 23:19:23.882819: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations: AVX2 AVX512F FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
Here we’ll use the GPflow functionality, but put a non-GPflow model inside the kernel.
Vanilla ConvNet. This gets 97.3% accuracy on MNIST when used on its own (+ final linear layer) after 20K iterations
[3]:
class KernelWithConvNN(gpflow.kernels.Kernel):
def __init__(
self,
image_shape: Tuple,
output_dim: int,
base_kernel: gpflow.kernels.Kernel,
batch_size: Optional[int] = None,
):
super().__init__()
with self.name_scope:
self.base_kernel = base_kernel
input_size = int(tf.reduce_prod(image_shape))
input_shape = (input_size,)
self.cnn = tf.keras.Sequential(
[
tf.keras.layers.InputLayer(
input_shape=input_shape, batch_size=batch_size
),
tf.keras.layers.Reshape(image_shape),
tf.keras.layers.Conv2D(
filters=32,
kernel_size=image_shape[:-1],
padding="same",
activation="relu",
),
tf.keras.layers.MaxPool2D(pool_size=(2, 2), strides=2),
tf.keras.layers.Conv2D(
filters=64,
kernel_size=(5, 5),
padding="same",
activation="relu",
),
tf.keras.layers.MaxPool2D(pool_size=(2, 2), strides=2),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(output_dim, activation="relu"),
tf.keras.layers.Lambda(to_default_float),
]
)
self.cnn.build()
def K(
self, a_input: tf.Tensor, b_input: Optional[tf.Tensor] = None
) -> tf.Tensor:
transformed_a = self.cnn(a_input)
transformed_b = self.cnn(b_input) if b_input is not None else b_input
return self.base_kernel.K(transformed_a, transformed_b)
def K_diag(self, a_input: tf.Tensor) -> tf.Tensor:
transformed_a = self.cnn(a_input)
return self.base_kernel.K_diag(transformed_a)
\(K_{uf}\) is in ConvNN output space, therefore we need to update Kuf multidispatch.
[4]:
class KernelSpaceInducingPoints(gpflow.inducing_variables.InducingPoints):
pass
@gpflow.covariances.Kuu.register(KernelSpaceInducingPoints, KernelWithConvNN)
def Kuu(inducing_variable, kernel, jitter=None):
func = gpflow.covariances.Kuu.dispatch(
gpflow.inducing_variables.InducingPoints, gpflow.kernels.Kernel
)
return func(inducing_variable, kernel.base_kernel, jitter=jitter)
@gpflow.covariances.Kuf.register(
KernelSpaceInducingPoints, KernelWithConvNN, object
)
def Kuf(inducing_variable, kernel, a_input):
return kernel.base_kernel(inducing_variable.Z, kernel.cnn(a_input))
Now we are ready to create and initialize the model:
[5]:
num_mnist_classes = 10
output_dim = 5
num_inducing_points = 100
images_subset, labels_subset = next(iter(dataset.batch(32)))
images_subset = tf.reshape(images_subset, [-1, image_size])
labels_subset = tf.reshape(labels_subset, [-1, 1])
kernel = KernelWithConvNN(
image_shape,
output_dim,
gpflow.kernels.SquaredExponential(),
batch_size=batch_size,
)
likelihood = gpflow.likelihoods.MultiClass(num_mnist_classes)
inducing_variable_kmeans = kmeans2(
images_subset.numpy(), num_inducing_points, minit="points"
)[0]
inducing_variable_cnn = kernel.cnn(inducing_variable_kmeans)
inducing_variable = KernelSpaceInducingPoints(inducing_variable_cnn)
model = gpflow.models.SVGP(
kernel,
likelihood,
inducing_variable=inducing_variable,
num_data=total_num_data,
num_latent_gps=num_mnist_classes,
)
2023-05-03 23:19:24.694760: W tensorflow/core/kernels/data/cache_dataset_ops.cc:856] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset will be discarded. This can happen if you have an input pipeline similar to `dataset.cache().take(k).repeat()`. You should use `dataset.take(k).cache().repeat()` instead.
And start optimization:
[6]:
data_iterator = iter(dataset)
adam_opt = tf.optimizers.Adam(0.001)
training_loss = model.training_loss_closure(data_iterator)
@tf.function
def optimization_step():
adam_opt.minimize(training_loss, var_list=model.trainable_variables)
for _ in range(iterations):
optimization_step()
Let’s do predictions after training. Don’t expect that we will get a good accuracy, because we haven’t run training for long enough.
[7]:
m, v = model.predict_y(images_subset)
preds = np.argmax(m, 1).reshape(labels_subset.numpy().shape)
correct = preds == labels_subset.numpy().astype(int)
acc = np.average(correct.astype(float)) * 100.0
print("Accuracy is {:.4f}%".format(acc))
2023-05-03 23:19:59.121190: W tensorflow/core/framework/cpu_allocator_impl.cc:82] Allocation of 102760448 exceeds 10% of free system memory.
2023-05-03 23:20:02.289735: W tensorflow/core/framework/cpu_allocator_impl.cc:82] Allocation of 102760448 exceeds 10% of free system memory.
2023-05-03 23:20:02.398769: W tensorflow/core/framework/cpu_allocator_impl.cc:82] Allocation of 102760448 exceeds 10% of free system memory.
2023-05-03 23:20:02.821898: W tensorflow/core/framework/cpu_allocator_impl.cc:82] Allocation of 102760448 exceeds 10% of free system memory.
2023-05-03 23:20:05.891315: W tensorflow/core/framework/cpu_allocator_impl.cc:82] Allocation of 102760448 exceeds 10% of free system memory.
Accuracy is 21.1914%