def test_random_feature_posterior_approximation(self): """Tests random feature GP's ability in approximating exact GP posterior.""" # Set momentum = 0.5 so posterior precision matrix is 0.5 * (I + K). gp_cov_momentum = 0.5 gp_cov_ridge_penalty = 1. num_inducing = 1024 rfgp_model = gaussian_process.RandomFeatureGaussianProcess( units=1, num_inducing=num_inducing, normalize_input=False, gp_kernel_type='gaussian', gp_cov_momentum=gp_cov_momentum, gp_cov_ridge_penalty=gp_cov_ridge_penalty) # Computes posterior covariance on test data. _, _ = rfgp_model(self.x_tr, training=True) _, gp_cov_ts = rfgp_model(self.x_ts, training=False) # Scale up covariance estimate since prec matrix is down-scaled by momentum. post_kernel_computed = gp_cov_ts * gp_cov_momentum post_kernel_expected = _compute_posterior_kernel(self.x_tr, self.x_ts, self.rbf_kern_func, gp_cov_ridge_penalty) np.testing.assert_allclose(post_kernel_computed, post_kernel_expected, **self.cov_tolerance)
def __init__(self, inner_dim, num_classes, cls_token_idx=0, activation="tanh", dropout_rate=0.0, initializer="glorot_uniform", use_spec_norm=True, use_gp_layer=True, temperature=None, **kwargs): """Initializes the `GaussianProcessClassificationHead`. Args: inner_dim: The dimensionality of inner projection layer. If 0 or `None` then only the output projection layer is created. num_classes: Number of output classes. cls_token_idx: The index inside the sequence to pool. activation: Dense layer activation. dropout_rate: Dropout probability. initializer: Initializer for dense layer kernels. use_spec_norm: Whether to apply spectral normalization to pooler layer. use_gp_layer: Whether to use Gaussian process as the output layer. temperature: The temperature parameter to be used for mean-field approximation during inference. If None then no mean-field adjustment is applied. **kwargs: Additional keyword arguments. """ # Collects spectral normalization and Gaussian process args from kwargs. self.use_spec_norm = use_spec_norm self.use_gp_layer = use_gp_layer self.spec_norm_kwargs = extract_spec_norm_kwargs(kwargs) self.gp_layer_kwargs = extract_gp_layer_kwargs(kwargs) self.temperature = temperature super().__init__(inner_dim=inner_dim, num_classes=num_classes, cls_token_idx=cls_token_idx, activation=activation, dropout_rate=dropout_rate, initializer=initializer, **kwargs) # Applies spectral normalization to the dense pooler layer. if self.use_spec_norm and hasattr(self, "dense"): self.dense = spectral_normalization.SpectralNormalization( self.dense, inhere_layer_name=True, **self.spec_norm_kwargs) # Replace Dense output layer with the Gaussian process layer. if use_gp_layer: self.out_proj = gaussian_process.RandomFeatureGaussianProcess( self.num_classes, kernel_initializer=tf_utils.clone_initializer( self.initializer), name="logits", **self.gp_layer_kwargs)
def __init__(self, inner_dim, num_classes, cls_token_idx=0, activation="tanh", dropout_rate=0.0, initializer="glorot_uniform", use_spec_norm=True, use_gp_layer=True, **kwargs): """Initializes the `GaussianProcessClassificationHead`. Args: inner_dim: The dimensionality of inner projection layer. num_classes: Number of output classes. cls_token_idx: The index inside the sequence to pool. activation: Dense layer activation. dropout_rate: Dropout probability. initializer: Initializer for dense layer kernels. use_spec_norm: Whether to apply spectral normalization to pooler layer. use_gp_layer: Whether to use Gaussian process as the output layer. **kwargs: Additional keyword arguments. """ # Collects spectral normalization and Gaussian process args from kwargs. self.use_spec_norm = use_spec_norm self.use_gp_layer = use_gp_layer self.spec_norm_kwargs = extract_spec_norm_kwargs(kwargs) self.gp_layer_kwargs = extract_gp_layer_kwargs(kwargs) super().__init__(inner_dim=inner_dim, num_classes=num_classes, cls_token_idx=cls_token_idx, activation=activation, dropout_rate=dropout_rate, initializer=initializer, **kwargs) # Applies spectral normalization to the pooler layer. if use_spec_norm: self.dense = spectral_normalization.SpectralNormalization( self.dense, inhere_layer_name=True, **self.spec_norm_kwargs) # Replace Dense output layer with the Gaussian process layer. if use_gp_layer: self.out_proj = gaussian_process.RandomFeatureGaussianProcess( self.num_classes, kernel_initializer=self.initializer, name="logits", **self.gp_layer_kwargs)
def test_no_matrix_update_during_test(self): """Tests if the precision matrix is not updated during testing.""" rfgp_model = gaussian_process.RandomFeatureGaussianProcess(units=1) # Training. _, gp_covmat_null = rfgp_model(self.x_tr, training=True) precision_mat_before_test = rfgp_model._gp_cov_layer.precision_matrix # Testing. _ = rfgp_model(self.x_ts, training=False) precision_mat_after_test = rfgp_model._gp_cov_layer.precision_matrix self.assertAllClose( gp_covmat_null, tf.eye(self.num_train_sample), atol=1e-4) self.assertAllClose( precision_mat_before_test, precision_mat_after_test, atol=1e-4)
def test_random_feature_prior_approximation(self): """Tests random feature GP's ability in approximating exact GP prior.""" num_inducing = 10240 rfgp_model = gaussian_process.RandomFeatureGaussianProcess( units=1, num_inducing=num_inducing, normalize_input=False, gp_kernel_type='gaussian', return_random_features=True) # Extract random features. _, _, gp_feature = rfgp_model(self.x_tr, training=True) gp_feature_np = gp_feature.numpy() prior_kernel_computed = gp_feature_np.dot(gp_feature_np.T) prior_kernel_expected = self.rbf_kern_func(self.x_tr, self.x_tr) np.testing.assert_allclose(prior_kernel_computed, prior_kernel_expected, **self.cov_tolerance)
def test_state_saving_and_loading(self): """Tests if the loaded model returns same results.""" input_data = np.random.random((1, 2)) rfgp_model = gaussian_process.RandomFeatureGaussianProcess(units=1) inputs = tf.keras.Input((2,), batch_size=1) outputs = rfgp_model(inputs) model = tf.keras.Model(inputs, outputs) gp_output, gp_covmat = model.predict(input_data) # Save and then load the model. temp_dir = self.get_temp_dir() self.addCleanup(shutil.rmtree, temp_dir) saved_model_dir = os.path.join(temp_dir, 'rfgp_model') model.save(saved_model_dir) new_model = tf.keras.models.load_model(saved_model_dir) gp_output_new, gp_covmat_new = new_model.predict(input_data) self.assertAllClose(gp_output, gp_output_new, atol=1e-4) self.assertAllClose(gp_covmat, gp_covmat_new, atol=1e-4)
def test_random_feature_linear_kernel(self): """Tests if linear kernel indeed leads to an identity mapping.""" # Specify linear kernel gp_kernel_type = 'linear' normalize_input = False scale_random_features = False use_custom_random_features = True rfgp_model = gaussian_process.RandomFeatureGaussianProcess( units=1, normalize_input=normalize_input, gp_kernel_type=gp_kernel_type, scale_random_features=scale_random_features, use_custom_random_features=use_custom_random_features, return_random_features=True) _, _, gp_feature = rfgp_model(self.x_tr, training=True) # Check if linear kernel leads to identity mapping. np.testing.assert_allclose(gp_feature, self.x_tr, **self.prec_tolerance)
def test_layer_build(self): """Tests if layer.built=True after building.""" rfgp_model = gaussian_process.RandomFeatureGaussianProcess(units=1) rfgp_model.build(input_shape=self.x_tr.shape) self.assertTrue(rfgp_model.built)