# Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Neural Network Gaussian Process (nngp) kernel computation.
Implementaion based on
"Deep Neural Networks as Gaussian Processes" by
Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S. Schoenholz,
Jeffrey Pennington, Jascha Sohl-Dickstein
arXiv:1711.00165 (https://arxiv.org/abs/1711.00165).
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import multiprocessing
import os
import numpy as np
import tensorflow as tf
import interp
flags = tf.app.flags
FLAGS = flags.FLAGS
flags.DEFINE_boolean("use_precomputed_grid", True,
"Option to save/load pre-computed grid")
flags.DEFINE_integer(
"fraction_of_int32", 32,
"allow batches at most of size int32.max / fraction_of_int32")
class NNGPKernel(object):
"""The iterative covariance Kernel for Neural Network Gaussian Process.
Args:
depth: int, number of hidden layers in corresponding NN.
nonlin_fn: tf ops corresponding to point-wise non-linearity in corresponding
NN. e.g.) tf.nn.relu, tf.nn.sigmoid, lambda x: x * tf.nn.sigmoid(x), ...
weight_var: initial value for the weight_variances parameter.
bias_var: initial value for the bias_variance parameter.
n_gauss: Number of gaussian integration grid. Choose odd integer, so that
there is a gridpoint at 0.
n_var: Number of variance grid points.
n_corr: Number of correlation grid points.
use_fixed_point_norm: bool, normalize input to variance fixed point.
Defaults to False, normalizing input to unit norm over input dimension.
"""
def __init__(self,
depth=1,
nonlin_fn=tf.tanh,
weight_var=1.,
bias_var=1.,
n_gauss=101,
n_var=151,
n_corr=131,
max_var=100,
max_gauss=100,
use_fixed_point_norm=False,
grid_path=None,
sess=None):
self.depth = depth
self.weight_var = weight_var
self.bias_var = bias_var
self.use_fixed_point_norm = use_fixed_point_norm
self.sess = sess
if FLAGS.use_precomputed_grid and (grid_path is None):
raise ValueError("grid_path must be specified to use precomputed grid.")
self.grid_path = grid_path
self.nonlin_fn = nonlin_fn
(self.var_aa_grid, self.corr_ab_grid, self.qaa_grid,
self.qab_grid) = self.get_grid(n_gauss, n_var, n_corr, max_var, max_gauss)
if self.use_fixed_point_norm:
self.var_fixed_point_np, self.var_fixed_point = self.get_var_fixed_point()
def get_grid(self, n_gauss, n_var, n_corr, max_var, max_gauss):
"""Get covariance grid by loading or computing a new one.
"""
# File configuration for precomputed grid
if FLAGS.use_precomputed_grid:
grid_path = self.grid_path
# TODO(jaehlee) np.save have broadcasting error when n_var==n_corr.
if n_var == n_corr:
n_var += 1
grid_file_name = "grid_{0:s}_ng{1:d}_ns{2:d}_nc{3:d}".format(
self.nonlin_fn.__name__, n_gauss, n_var, n_corr)
grid_file_name += "_mv{0:d}_mg{1:d}".format(max_var, max_gauss)
# Load grid file if it exists already
if (FLAGS.use_precomputed_grid and
tf.gfile.Exists(os.path.join(grid_path, grid_file_name))):
with tf.gfile.Open(os.path.join(grid_path, grid_file_name), "rb") as f:
grid_data_np = np.load(f)
tf.logging.info("Loaded interpolation grid from %s"%
os.path.join(grid_path, grid_file_name))
grid_data = (tf.convert_to_tensor(grid_data_np[0], dtype=tf.float64),
tf.convert_to_tensor(grid_data_np[1], dtype=tf.float64),
tf.convert_to_tensor(grid_data_np[2], dtype=tf.float64),
tf.convert_to_tensor(grid_data_np[3], dtype=tf.float64))
else:
tf.logging.info("Generating interpolation grid...")
grid_data = _compute_qmap_grid(self.nonlin_fn, n_gauss, n_var, n_corr,
max_var=max_var, max_gauss=max_gauss)
if FLAGS.use_precomputed_grid:
with tf.Session() as sess:
grid_data_np = sess.run(grid_data)
tf.gfile.MakeDirs(grid_path)
with tf.gfile.Open(os.path.join(grid_path, grid_file_name), "wb") as f:
np.save(f, grid_data_np)
with tf.gfile.Open(os.path.join(grid_path, grid_file_name), "rb") as f:
grid_data_np = np.load(f)
tf.logging.info("Loaded interpolation grid from %s"%
os.path.join(grid_path, grid_file_name))
grid_data = (tf.convert_to_tensor(grid_data_np[0], dtype=tf.float64),
tf.convert_to_tensor(grid_data_np[1], dtype=tf.float64),
tf.convert_to_tensor(grid_data_np[2], dtype=tf.float64),
tf.convert_to_tensor(grid_data_np[3], dtype=tf.float64))
return grid_data
def get_var_fixed_point(self):
with tf.name_scope("get_var_fixed_point"):
# If normalized input length starts at 1.
current_qaa = self.weight_var * tf.constant(
[1.], dtype=tf.float64) + self.bias_var
diff = 1.
prev_qaa_np = 1.
it = 0
while diff > 1e-6 and it < 300:
samp_qaa = interp.interp_lin(
self.var_aa_grid, self.qaa_grid, current_qaa)
samp_qaa = self.weight_var * samp_qaa + self.bias_var
current_qaa = samp_qaa
with tf.Session() as sess:
current_qaa_np = sess.run(current_qaa)
diff = np.abs(current_qaa_np - prev_qaa_np)
it += 1
prev_qaa_np = current_qaa_np
return current_qaa_np, current_qaa
def k_diag(self, input_x, return_full=True):
"""Iteratively building the diagonal part (variance) of the NNGP kernel.
Args:
input_x: tensor of input of size [num_data, input_dim].
return_full: boolean for output to be [num_data] sized or a scalar value
for normalized inputs
Sets self.layer_qaa_dict of {layer #: qaa at the layer}
Returns:
qaa: variance at the output.
"""
with tf.name_scope("Kdiag"):
# If normalized input length starts at 1.
if self.use_fixed_point_norm:
current_qaa = self.var_fixed_point
else:
current_qaa = self.weight_var * tf.convert_to_tensor(
[1.], dtype=tf.float64) + self.bias_var
self.layer_qaa_dict = {0: current_qaa}
for l in xrange(self.depth):
with tf.name_scope("layer_%d" % l):
samp_qaa = interp.interp_lin(
self.var_aa_grid, self.qaa_grid, current_qaa)
samp_qaa = self.weight_var * samp_qaa + self.bias_var
self.layer_qaa_dict[l + 1] = samp_qaa
current_qaa = samp_qaa
if return_full:
qaa = tf.tile(current_qaa[:1], ([input_x.shape[0].value]))
else:
qaa = current_qaa[0]
return qaa
def k_full(self, input1, input2=None):
"""Iteratively building the full NNGP kernel.
"""
input1 = self._input_layer_normalization(input1)
if input2 is None:
input2 = input1
else:
input2 = self._input_layer_normalization(input2)
with tf.name_scope("k_full"):
cov_init = tf.matmul(
input1, input2, transpose_b=True) / input1.shape[1].value
self.k_diag(input1)
q_aa_init = self.layer_qaa_dict[0]
q_ab = cov_init
q_ab = self.weight_var * q_ab + self.bias_var
corr = q_ab / q_aa_init[0]
if FLAGS.fraction_of_int32 > 1:
batch_size, batch_count = self._get_batch_size_and_count(input1, input2)
with tf.name_scope("q_ab"):
q_ab_all = []
for b_x in range(batch_count):
with tf.name_scope("batch_%d" % b_x):
corr_flat_batch = corr[
batch_size * b_x : batch_size * (b_x + 1), :]
corr_flat_batch = tf.reshape(corr_flat_batch, [-1])
for l in xrange(self.depth):
with tf.name_scope("layer_%d" % l):
q_aa = self.layer_qaa_dict[l]
q_ab = interp.interp_lin_2d(x=self.var_aa_grid,
y=self.corr_ab_grid,
z=self.qab_grid,
xp=q_aa,
yp=corr_flat_batch)
q_ab = self.weight_var * q_ab + self.bias_var
corr_flat_batch = q_ab / self.layer_qaa_dict[l + 1][0]
q_ab_all.append(q_ab)
q_ab_all = tf.parallel_stack(q_ab_all)
else:
with tf.name_scope("q_ab"):
corr_flat = tf.reshape(corr, [-1])
for l in xrange(self.depth):
with tf.name_scope("layer_%d" % l):
q_aa = self.layer_qaa_dict[l]
q_ab = interp.interp_lin_2d(x=self.var_aa_grid,
y=self.corr_ab_grid,
z=self.qab_grid,
xp=q_aa,
yp=corr_flat)
q_ab = self.weight_var * q_ab + self.bias_var
corr_flat = q_ab / self.layer_qaa_dict[l+1][0]
q_ab_all = q_ab
return tf.reshape(q_ab_all, cov_init.shape, "qab")
def _input_layer_normalization(self, x):
"""Input normalization to unit variance or fixed point variance.
"""
with tf.name_scope("input_layer_normalization"):
# Layer norm, fix to unit variance
eps = 1e-15
mean, var = tf.nn.moments(x, axes=[1], keep_dims=True)
x_normalized = (x - mean) / tf.sqrt(var + eps)
if self.use_fixed_point_norm:
x_normalized *= tf.sqrt(
(self.var_fixed_point[0] - self.bias_var) / self.weight_var)
return x_normalized
def _get_batch_size_and_count(self, input1, input2):
"""Compute batch size and number to split when input size is large.
Args:
input1: tensor, input tensor to covariance matrix
input2: tensor, second input tensor to covariance matrix
Returns:
batch_size: int, size of each batch
batch_count: int, number of batches
"""
input1_size = input1.shape[0].value
input2_size = input2.shape[0].value
batch_size = min(np.iinfo(np.int32).max //
(FLAGS.fraction_of_int32 * input2_size), input1_size)
while input1_size % batch_size != 0:
batch_size -= 1
batch_count = input1_size // batch_size
return batch_size, batch_count
def _fill_qab_slice(idx, z1, z2, var_aa, corr_ab, nonlin_fn):
"""Helper method used for parallel computation for full qab."""
log_weights_ab_unnorm = -(z1**2 + z2**2 - 2 * z1 * z2 * corr_ab) / (
2 * var_aa[idx] * (1 - corr_ab**2))
log_weights_ab = log_weights_ab_unnorm - tf.reduce_logsumexp(
log_weights_ab_unnorm, axis=[0, 1], keep_dims=True)
weights_ab = tf.exp(log_weights_ab)
qab_slice = tf.reduce_sum(
nonlin_fn(z1) * nonlin_fn(z2) * weights_ab, axis=[0, 1])
qab_slice = tf.Print(qab_slice, [idx], "Generating slice: ")
return qab_slice
def _compute_qmap_grid(nonlin_fn,
n_gauss,
n_var,
n_corr,
log_spacing=False,
min_var=1e-8,
max_var=100.,
max_corr=0.99999,
max_gauss=10.):
"""Construct graph for covariance grid to use for kernel computation.
Given variance and correlation (or covariance) of pre-activation, perform
Gaussian integration to get covariance of post-activation.
Raises:
ValueError: if n_gauss is even integer.
Args:
nonlin_fn: tf ops corresponding to point-wise non-linearity in
corresponding NN. e.g.) tf.nn.relu, tf.nn.sigmoid,
lambda x: x * tf.nn.sigmoid(x), ...
n_gauss: int, number of Gaussian integration points with equal spacing
between (-max_gauss, max_gauss). Choose odd integer, so that there is a
gridpoint at 0.
n_var: int, number of variance grid points.get_grid
n_corr: int, number of correlation grid points.
log_spacing: bool, whether to use log-linear instead of linear variance
grid.
min_var: float, smallest variance value to generate grid.
max_var: float, largest varaince value to generate grid.
max_corr: float, largest correlation value to generate grid. Should be
slightly smaller than 1.
max_gauss: float, range (-max_gauss, max_gauss) for Gaussian integration.
Returns:
var_grid_pts: tensor of size [n_var], grid points where variance are
evaluated at.
corr_grid_pts: tensor of size [n_corr], grid points where correlation are
evalutated at.
qaa: tensor of size [n_var], variance of post-activation at given
pre-activation variance.
qab: tensor of size [n_var, n_corr], covariance of post-activation at
given pre-activation variance and correlation.
"""
if n_gauss % 2 != 1:
raise ValueError("n_gauss=%d should be an odd integer" % n_gauss)
with tf.name_scope("compute_qmap_grid"):
min_var = tf.convert_to_tensor(min_var, dtype=tf.float64)
max_var = tf.convert_to_tensor(max_var, dtype=tf.float64)
max_corr = tf.convert_to_tensor(max_corr, dtype=tf.float64)
max_gauss = tf.convert_to_tensor(max_gauss, dtype=tf.float64)
# Evaluation points for numerical integration over a Gaussian.
z1 = tf.reshape(tf.linspace(-max_gauss, max_gauss, n_gauss), (-1, 1, 1))
z2 = tf.transpose(z1, perm=[1, 0, 2])
if log_spacing:
var_aa = tf.exp(tf.linspace(tf.log(min_var), tf.log(max_var), n_var))
else:
# Evaluation points for pre-activations variance and correlation
var_aa = tf.linspace(min_var, max_var, n_var)
corr_ab = tf.reshape(tf.linspace(-max_corr, max_corr, n_corr), (1, 1, -1))
# compute q_aa
log_weights_aa_unnorm = -0.5 * (z1**2 / tf.reshape(var_aa, [1, 1, -1]))
log_weights_aa = log_weights_aa_unnorm - tf.reduce_logsumexp(
log_weights_aa_unnorm, axis=[0, 1], keep_dims=True)
weights_aa = tf.exp(log_weights_aa)
qaa = tf.reduce_sum(nonlin_fn(z1)**2 * weights_aa, axis=[0, 1])
# compute q_ab
# weights to reweight uniform samples by, for q_ab.
# (weights are probability of z1, z2 under Gaussian
# w/ variance var_aa and covariance var_aa*corr_ab)
# weights_ab will have shape [n_g, n_g, n_v, n_c]
def fill_qab_slice(idx):
return _fill_qab_slice(idx, z1, z2, var_aa, corr_ab, nonlin_fn)
qab = tf.map_fn(
fill_qab_slice,
tf.range(n_var),
dtype=tf.float64,
parallel_iterations=multiprocessing.cpu_count())
var_grid_pts = tf.reshape(var_aa, [-1])
corr_grid_pts = tf.reshape(corr_ab, [-1])
return var_grid_pts, corr_grid_pts, qaa, qab
