From b4b47b1403da433ab5f45c1b27eed3e79b727949 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 28 Aug 2023 07:50:20 -0700 Subject: [PATCH] Generalize the testing API to support input Tensors of dimension >1, excluding the batch dimension. This is a forward-looking change for testing more general layers such as `tf.keras.layers.LayerNormalization` and `tf.keras.layers.EinsumDense`. PiperOrigin-RevId: 560709678 --- .../fast_gradient_clipping/clip_grads_test.py | 10 +- .../common_test_utils.py | 120 +++++++++++------- .../registry_functions/dense_test.py | 30 +++-- .../registry_functions/embedding_test.py | 2 +- .../fast_gradient_clipping/type_aliases.py | 2 +- 5 files changed, 95 insertions(+), 69 deletions(-) diff --git a/tensorflow_privacy/privacy/fast_gradient_clipping/clip_grads_test.py b/tensorflow_privacy/privacy/fast_gradient_clipping/clip_grads_test.py index 2e12a1c..5f6eda2 100644 --- a/tensorflow_privacy/privacy/fast_gradient_clipping/clip_grads_test.py +++ b/tensorflow_privacy/privacy/fast_gradient_clipping/clip_grads_test.py @@ -110,10 +110,10 @@ class CustomLayerTest(tf.test.TestCase, parameterized.TestCase): continue (computed_norms, true_norms) = ( common_test_utils.get_computed_and_true_norms( - model_generator=common_test_utils.make_two_layer_sequential_model, - layer_generator=lambda a, b: DoubleDense(b), - input_dims=input_dim, - output_dim=output_dim, + model_generator=common_test_utils.make_two_layer_functional_model, + layer_generator=lambda a, b: DoubleDense(*b), + input_dims=[input_dim], + output_dims=[output_dim], per_example_loss_fn=per_example_loss_fn, num_microbatches=num_microbatches, is_eager=is_eager, @@ -136,7 +136,7 @@ class ComputeClippedGradsAndOutputsTest( dense_generator = lambda a, b: tf.keras.layers.Dense(b) self._input_dim = 2 self._output_dim = 3 - self._model = common_test_utils.make_two_layer_sequential_model( + self._model = common_test_utils.make_two_layer_functional_model( dense_generator, self._input_dim, self._output_dim ) diff --git a/tensorflow_privacy/privacy/fast_gradient_clipping/common_test_utils.py b/tensorflow_privacy/privacy/fast_gradient_clipping/common_test_utils.py index e99698b..e5f5e09 100644 --- a/tensorflow_privacy/privacy/fast_gradient_clipping/common_test_utils.py +++ b/tensorflow_privacy/privacy/fast_gradient_clipping/common_test_utils.py @@ -13,7 +13,7 @@ # limitations under the License. """A collection of common utility functions for unit testing.""" -from typing import Callable, List, Optional, Tuple, Union +from typing import Callable, List, Optional, Tuple import numpy as np import tensorflow as tf @@ -108,12 +108,12 @@ def compute_true_gradient_norms( def get_model_from_generator( model_generator: type_aliases.ModelGenerator, layer_generator: type_aliases.LayerGenerator, - input_dims: Union[int, List[int]], - output_dim: int, + input_dims: List[int], + output_dims: List[int], is_eager: bool, ) -> tf.keras.Model: """Creates a simple model from input specifications.""" - model = model_generator(layer_generator, input_dims, output_dim) + model = model_generator(layer_generator, input_dims, output_dims) model.compile( optimizer=tf.keras.optimizers.SGD(learning_rate=1.0), loss=tf.keras.losses.MeanSquaredError( @@ -171,8 +171,8 @@ def get_computed_and_true_norms_from_model( def get_computed_and_true_norms( model_generator: type_aliases.ModelGenerator, layer_generator: type_aliases.LayerGenerator, - input_dims: Union[int, List[int]], - output_dim: int, + input_dims: List[int], + output_dims: List[int], per_example_loss_fn: Optional[Callable[[tf.Tensor, tf.Tensor], tf.Tensor]], num_microbatches: Optional[int], is_eager: bool, @@ -196,7 +196,7 @@ def get_computed_and_true_norms( Returns a `tf.keras.layers.Layer` that accepts input tensors of dimension `idim` and returns output tensors of dimension `odim`. input_dims: The input dimension(s) of the test `tf.keras.Model` instance. - output_dim: The output dimension of the test `tf.keras.Model` instance. + output_dims: The output dimension(s) of the test `tf.keras.Model` instance. per_example_loss_fn: If not None, used as vectorized per example loss function. num_microbatches: The number of microbatches. None or an integer. @@ -219,7 +219,7 @@ def get_computed_and_true_norms( model_generator=model_generator, layer_generator=layer_generator, input_dims=input_dims, - output_dim=output_dim, + output_dims=output_dims, is_eager=is_eager, ) return get_computed_and_true_norms_from_model( @@ -234,64 +234,74 @@ def get_computed_and_true_norms( ) +def reshape_and_sum(tensor: tf.Tensor) -> tf.Tensor: + """Reshapes and sums along non-batch dims to get the shape [None, 1].""" + reshaped_2d = tf.reshape(tensor, [tf.shape(tensor)[0], -1]) + return tf.reduce_sum(reshaped_2d, axis=-1, keepdims=True) + + # ============================================================================== # Model generators. # ============================================================================== -def make_two_layer_sequential_model(layer_generator, input_dim, output_dim): - """Creates a 2-layer sequential model.""" - model = tf.keras.Sequential() - model.add(tf.keras.Input(shape=(input_dim,))) - model.add(layer_generator(input_dim, output_dim)) - model.add(tf.keras.layers.Dense(1)) - return model - - -def make_three_layer_sequential_model(layer_generator, input_dim, output_dim): - """Creates a 3-layer sequential model.""" - model = tf.keras.Sequential() - model.add(tf.keras.Input(shape=(input_dim,))) - layer1 = layer_generator(input_dim, output_dim) - model.add(layer1) - if isinstance(layer1, tf.keras.layers.Embedding): - # Having multiple consecutive embedding layers does not make sense since - # embedding layers only map integers to real-valued vectors. - model.add(tf.keras.layers.Dense(output_dim)) - else: - model.add(layer_generator(output_dim, output_dim)) - model.add(tf.keras.layers.Dense(1)) - return model - - -def make_two_layer_functional_model(layer_generator, input_dim, output_dim): - """Creates a 2-layer 1-input functional model with a pre-output square op.""" - inputs = tf.keras.Input(shape=(input_dim,)) - layer1 = layer_generator(input_dim, output_dim) +def make_one_layer_functional_model( + layer_generator: type_aliases.LayerGenerator, + input_dims: List[int], + output_dims: List[int], +) -> tf.keras.Model: + """Creates a 1-layer sequential model.""" + inputs = tf.keras.Input(shape=input_dims) + layer1 = layer_generator(input_dims, output_dims) temp1 = layer1(inputs) - temp2 = tf.square(temp1) - outputs = tf.keras.layers.Dense(1)(temp2) + outputs = reshape_and_sum(temp1) return tf.keras.Model(inputs=inputs, outputs=outputs) -def make_two_tower_model(layer_generator, input_dim, output_dim): +def make_two_layer_functional_model( + layer_generator: type_aliases.LayerGenerator, + input_dims: List[int], + output_dims: List[int], +) -> tf.keras.Model: + """Creates a 2-layer sequential model.""" + inputs = tf.keras.Input(shape=input_dims) + layer1 = layer_generator(input_dims, output_dims) + temp1 = layer1(inputs) + temp2 = tf.keras.layers.Dense(1)(temp1) + outputs = reshape_and_sum(temp2) + return tf.keras.Model(inputs=inputs, outputs=outputs) + + +def make_two_tower_model( + layer_generator: type_aliases.LayerGenerator, + input_dims: List[int], + output_dims: List[int], +) -> tf.keras.Model: """Creates a 2-layer 2-input functional model.""" - inputs1 = tf.keras.Input(shape=(input_dim,)) - layer1 = layer_generator(input_dim, output_dim) + inputs1 = tf.keras.Input(shape=input_dims) + layer1 = layer_generator(input_dims, output_dims) temp1 = layer1(inputs1) - inputs2 = tf.keras.Input(shape=(input_dim,)) - layer2 = layer_generator(input_dim, output_dim) + inputs2 = tf.keras.Input(shape=input_dims) + layer2 = layer_generator(input_dims, output_dims) temp2 = layer2(inputs2) temp3 = tf.add(temp1, temp2) - outputs = tf.keras.layers.Dense(1)(temp3) + temp4 = tf.keras.layers.Dense(1)(temp3) + outputs = reshape_and_sum(temp4) return tf.keras.Model(inputs=[inputs1, inputs2], outputs=outputs) -def make_bow_model(layer_generator, input_dims, output_dim): +def make_bow_model( + layer_generator: type_aliases.LayerGenerator, + input_dims: List[int], + output_dims: List[int], +) -> tf.keras.Model: """Creates a simple embedding bow model.""" del layer_generator inputs = tf.keras.Input(shape=input_dims) # For the Embedding layer, input_dim is the vocabulary size. This should # be distinguished from the input_dim argument, which is the number of ids # in eache example. + if len(output_dims) != 1: + raise ValueError('Expected `output_dims` to be of size 1.') + output_dim = output_dims[0] emb_layer = tf.keras.layers.Embedding(input_dim=10, output_dim=output_dim) feature_embs = emb_layer(inputs) # Embeddings add one extra dimension to its inputs, which combined with the @@ -305,7 +315,11 @@ def make_bow_model(layer_generator, input_dims, output_dim): return tf.keras.Model(inputs=inputs, outputs=example_embs) -def make_dense_bow_model(layer_generator, input_dims, output_dim): +def make_dense_bow_model( + layer_generator: type_aliases.LayerGenerator, + input_dims: List[int], + output_dims: List[int], +) -> tf.keras.Model: """Creates an embedding bow model with a `Dense` layer.""" del layer_generator inputs = tf.keras.Input(shape=input_dims) @@ -313,6 +327,9 @@ def make_dense_bow_model(layer_generator, input_dims, output_dim): # be distinguished from the input_dim argument, which is the number of ids # in eache example. cardinality = 10 + if len(output_dims) != 1: + raise ValueError('Expected `output_dims` to be of size 1.') + output_dim = output_dims[0] emb_layer = tf.keras.layers.Embedding( input_dim=cardinality, output_dim=output_dim ) @@ -329,7 +346,11 @@ def make_dense_bow_model(layer_generator, input_dims, output_dim): return tf.keras.Model(inputs=inputs, outputs=outputs) -def make_weighted_bow_model(layer_generator, input_dims, output_dim): +def make_weighted_bow_model( + layer_generator: type_aliases.LayerGenerator, + input_dims: List[int], + output_dims: List[int], +) -> tf.keras.Model: """Creates a weighted embedding bow model.""" # NOTE: This model only accepts dense input tensors. del layer_generator @@ -338,6 +359,9 @@ def make_weighted_bow_model(layer_generator, input_dims, output_dim): # be distinguished from the input_dim argument, which is the number of ids # in eache example. cardinality = 10 + if len(output_dims) != 1: + raise ValueError('Expected `output_dims` to be of size 1.') + output_dim = output_dims[0] emb_layer = tf.keras.layers.Embedding( input_dim=cardinality, output_dim=output_dim ) diff --git a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/dense_test.py b/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/dense_test.py index c0b7f7b..3bde1ad 100644 --- a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/dense_test.py +++ b/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/dense_test.py @@ -23,21 +23,21 @@ from tensorflow_privacy.privacy.fast_gradient_clipping.registry_functions import # Helper functions. # ============================================================================== def get_dense_layer_generators(): - def sigmoid_dense_layer(b): - return tf.keras.layers.Dense(b, activation='sigmoid') + + def sigmoid_dense_layer(units): + return tf.keras.layers.Dense(units, activation='sigmoid') return { - 'pure_dense': lambda a, b: tf.keras.layers.Dense(b), - 'sigmoid_dense': lambda a, b: sigmoid_dense_layer(b), + 'pure_dense': lambda a, b: tf.keras.layers.Dense(b[0]), + 'sigmoid_dense': lambda a, b: sigmoid_dense_layer(b[0]), } def get_dense_model_generators(): return { - 'seq1': common_test_utils.make_two_layer_sequential_model, - 'seq2': common_test_utils.make_three_layer_sequential_model, - 'func1': common_test_utils.make_two_layer_functional_model, - 'tower1': common_test_utils.make_two_tower_model, + 'func1': common_test_utils.make_one_layer_functional_model, + 'func2': common_test_utils.make_two_layer_functional_model, + 'tower2': common_test_utils.make_two_tower_model, } @@ -62,7 +62,7 @@ class GradNormTest(tf.test.TestCase, parameterized.TestCase): @parameterized.product( model_name=list(get_dense_model_generators().keys()), layer_name=list(get_dense_layer_generators().keys()), - input_dim=[4], + input_dims=[[4]], output_dim=[2], layer_registry_name=list(get_dense_layer_registries().keys()), per_example_loss_fn=[None, common_test_utils.test_loss_fn], @@ -75,7 +75,7 @@ class GradNormTest(tf.test.TestCase, parameterized.TestCase): self, model_name, layer_name, - input_dim, + input_dims, output_dim, layer_registry_name, per_example_loss_fn, @@ -85,15 +85,17 @@ class GradNormTest(tf.test.TestCase, parameterized.TestCase): weighted, ): # Parse inputs to generate test data. - x_batches, weight_batches = common_test_utils.get_nd_test_batches(input_dim) + x_batches, weight_batches = common_test_utils.get_nd_test_batches( + input_dims[0] + ) # Load shared assets to all devices. with self.strategy.scope(): model = common_test_utils.get_model_from_generator( model_generator=get_dense_model_generators()[model_name], layer_generator=get_dense_layer_generators()[layer_name], - input_dims=input_dim, - output_dim=output_dim, + input_dims=input_dims, + output_dims=[output_dim], is_eager=is_eager, ) @@ -103,7 +105,7 @@ class GradNormTest(tf.test.TestCase, parameterized.TestCase): model=model, per_example_loss_fn=per_example_loss_fn, num_microbatches=num_microbatches, - x_batch=[x_batch, x_batch] if model_name == 'tower1' else x_batch, + x_batch=[x_batch, x_batch] if model_name == 'tower2' else x_batch, weight_batch=weight_batch if weighted else None, registry=get_dense_layer_registries()[layer_registry_name], partial=partial, diff --git a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/embedding_test.py b/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/embedding_test.py index f8c1cf3..44de2e4 100644 --- a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/embedding_test.py +++ b/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/embedding_test.py @@ -119,7 +119,7 @@ class GradNormTest(tf.test.TestCase, parameterized.TestCase): model_generator=get_embedding_model_generators()[model_name], layer_generator=None, input_dims=embed_indices.shape[1:], - output_dim=output_dim, + output_dims=[output_dim], is_eager=is_eager, ) diff --git a/tensorflow_privacy/privacy/fast_gradient_clipping/type_aliases.py b/tensorflow_privacy/privacy/fast_gradient_clipping/type_aliases.py index e8b7f91..181da5a 100644 --- a/tensorflow_privacy/privacy/fast_gradient_clipping/type_aliases.py +++ b/tensorflow_privacy/privacy/fast_gradient_clipping/type_aliases.py @@ -45,5 +45,5 @@ GeneratorFunction = Optional[Callable[[Any, Tuple, Dict], Tuple[Any, Any]]] LayerGenerator = Callable[[int, int], tf.keras.layers.Layer] ModelGenerator = Callable[ - [LayerGenerator, Union[int, List[int]], int], tf.keras.Model + [LayerGenerator, List[int], List[int]], tf.keras.Model ]