def test_saved_model_tf(self, adapted): input_data = [[0.0], [2.0], [0.0], [2.0]] expected_output = [[-1.0], [1.0], [-1.0], [1.0]] inputs = keras.Input(shape=(1,), dtype=tf.float32) if adapted: layer = normalization.Normalization(axis=-1) layer.adapt(input_data) else: layer = normalization.Normalization(mean=1.0, variance=1.0) outputs = layer(inputs) model = keras.Model(inputs=inputs, outputs=outputs) output_data = model.predict(input_data) self.assertAllClose(output_data, expected_output) # Save the model to disk. output_path = os.path.join(self.get_temp_dir(), "tf_saved_model") tf.saved_model.save(model, output_path) loaded_model = tf.saved_model.load(output_path) f = loaded_model.signatures["serving_default"] # Ensure that the loaded model is unique (so that the save/load is real) self.assertIsNot(model, loaded_model) # Validate correctness of the new model. new_output_data = f(tf.constant(input_data))["normalization"] self.assertAllClose(new_output_data, expected_output)
def test_layer_computation(self, adapt_data, axis, test_data, use_dataset, expected): input_shape = tuple( [test_data.shape[i] for i in range(1, test_data.ndim)]) if use_dataset: # Keras APIs expect batched datasets adapt_data = tf.data.Dataset.from_tensor_slices(adapt_data).batch( test_data.shape[0] // 2) test_data = tf.data.Dataset.from_tensor_slices(test_data).batch( test_data.shape[0] // 2) layer = normalization.Normalization(axis=axis) layer.adapt(adapt_data) input_data = keras.Input(shape=input_shape) output = layer(input_data) model = keras.Model(input_data, output) model._run_eagerly = testing_utils.should_run_eagerly() output_data = model.predict(test_data) self.assertAllClose(expected, output_data) weights = layer.get_weights() mean = weights[0] var = weights[1] direct_set_layer = normalization.Normalization(axis=axis, mean=mean, variance=var) input_data = keras.Input(shape=input_shape) output = direct_set_layer(input_data) model = keras.Model(input_data, output) model._run_eagerly = testing_utils.should_run_eagerly() output_data = model.predict(test_data) self.assertAllClose(expected, output_data)
def test_merge_state(self): if not tf.executing_eagerly(): self.skipTest("`merge_state` only supported in TF2") data = np.random.rand(30, 10, 2) ds = tf.data.Dataset.from_tensor_slices(data).batch(2) norm = normalization.Normalization(axis=(1, 2)) norm.adapt(ds) partial_ds_1 = ds.shard(3, 0) partial_ds_2 = ds.shard(3, 1) partial_ds_3 = ds.shard(3, 2) norm_1 = normalization.Normalization(axis=(1, 2)) norm_2 = normalization.Normalization(axis=(1, 2)) norm_3 = normalization.Normalization(axis=(1, 2)) norm_1.adapt(partial_ds_1) norm_2.adapt(partial_ds_2) norm_3.adapt(partial_ds_3) norm_1.merge_state([norm_2, norm_3]) merged_norm = norm_1 self.assertAllClose(norm(data), merged_norm(data))
def test_invert(self): data = np.array([0.0, 2.0, 0.0, 2.0]) norm = normalization.Normalization(mean=1.0, variance=1.0) inv_norm = normalization.Normalization( mean=1.0, variance=1.0, invert=True ) output = norm(data) output2 = inv_norm(output) self.assertListEqual(output2.shape.as_list(), [4]) self.assertAllClose(output2, [0.0, 2.0, 0.0, 2.0])
def test_output_dtype(self): if not tf.__internal__.tf2.enabled(): self.skipTest("set_global_policy only supported in TF2.") # Output should respect an explicit dtype, and default to the global policy. policy.set_global_policy("float64") input_data = keras.Input(batch_size=16, shape=(1,)) layer = normalization.Normalization(mean=1.0, variance=1.0, dtype="float16") output = layer(input_data) self.assertAllEqual(output.dtype, tf.float16) layer = normalization.Normalization(mean=1.0, variance=1.0) output = layer(input_data) self.assertAllEqual(output.dtype, tf.float64)
def run_dataset_implementation(self, num_elements, batch_size): input_t = keras.Input(shape=(1, )) layer = normalization.Normalization() _ = layer(input_t) num_repeats = 5 starts = [] ends = [] for _ in range(num_repeats): ds = tf.data.Dataset.range(num_elements) ds = ds.map( lambda x: tf.compat.v1.expand_dims(tf.cast(x, tf.float32), -1)) ds = ds.batch(batch_size) starts.append(time.time()) # Benchmarked code begins here. k, n, ex, ex2 = ds.reduce((0.0, 0, 0.0, 0.0), reduce_fn) mean = k.numpy() + ex.numpy() / n.numpy() var = (ex2.numpy() - (ex.numpy() * ex.numpy()) / n.numpy()) / (n.numpy() - 1) layer.set_weights([mean, var]) # Benchmarked code ends here. ends.append(time.time()) avg_time = np.mean(np.array(ends) - np.array(starts)) return avg_time
def bm_adapt_implementation(self, num_elements, batch_size): """Test the KPL adapt implementation.""" input_t = keras.Input(shape=(1, ), dtype=tf.float32) layer = normalization.Normalization() _ = layer(input_t) num_repeats = 5 starts = [] ends = [] for _ in range(num_repeats): ds = tf.data.Dataset.range(num_elements) ds = ds.map( lambda x: tf.compat.v1.expand_dims(tf.cast(x, tf.float32), -1)) ds = ds.batch(batch_size) starts.append(time.time()) # Benchmarked code begins here. layer.adapt(ds) # Benchmarked code ends here. ends.append(time.time()) avg_time = np.mean(np.array(ends) - np.array(starts)) name = "normalization_adapt|%s_elements|batch_%s" % (num_elements, batch_size) baseline = self.run_dataset_implementation(num_elements, batch_size) extras = { "tf.data implementation baseline": baseline, "delta seconds": (baseline - avg_time), "delta percent": ((baseline - avg_time) / baseline) * 100 } self.report_benchmark(iters=num_repeats, wall_time=avg_time, extras=extras, name=name)
def test_0d_unbatched_adapt(self): ds = tf.data.Dataset.from_tensor_slices([2., 0., 2., 0.]) layer = normalization.Normalization(axis=None) layer.adapt(ds) output_ds = ds.map(layer) self.assertAllClose(list(output_ds.as_numpy_iterator()), [1., -1., 1., -1.])
def test_list_input(self): with self.assertRaisesRegex( ValueError, ("Normalization only accepts a single input. If you are " "passing a python list or tuple as a single input, " "please convert to a numpy array or `tf.Tensor`.")): normalization.Normalization()([1, 2, 3])
def test_1d_data(self): data = [0, 2, 0, 2] layer = normalization.Normalization(axis=-1) layer.adapt(data) output = layer(data) self.assertListEqual(output.shape.as_list(), [4, 1]) if tf.executing_eagerly(): self.assertAllClose(output.numpy(), [[-1], [1], [-1], [1]])
def test_broadcasting_during_direct_setting(self): layer = normalization.Normalization(axis=-1, mean=[1.0], variance=[1.0]) output = layer(np.array([[1., 2.]])) expected_output = [[0., 1.]] self.assertAllClose(output, expected_output) self.assertAllClose(layer.get_weights(), [])
def test_broadcasting_during_direct_setting_with_tensors(self): layer = normalization.Normalization(axis=-1, mean=tf.constant([1.0]), variance=tf.constant([2.0])) layer.build((None, 2)) weights = layer.get_weights() self.assertAllClose([1.0, 1.0], weights[0]) self.assertAllClose([2.0, 2.0], weights[1])
def test_axis_permutations(self, axis): layer = normalization.Normalization(axis=axis) # data.shape = [2, 2, 3] data = np.array([[[0., 1., 2.], [0., 2., 6.]], [[2., 3., 4.], [3., 6., 10.]]]) expect = np.array([[[-1., -1., -1.], [-1., -1., -1.]], [[1., 1., 1.], [1., 1., 1.]]]) layer.adapt(data) self.assertAllClose(expect, layer(data))
def test_model_summary_after_layer_adapt(self): data = np.array([[[0., 1., 2.], [0., 2., 6.]], [[2., 3., 4.], [3., 6., 10.]]]) layer = normalization.Normalization(axis=-1) layer.adapt(data) model = keras.Sequential( [layer, keras.layers.Dense(64, activation="relu"), keras.layers.Dense(1)]) model.summary()
def test_0d_data(self): if not tf.executing_eagerly(): self.skipTest("Only supported in TF2.") data = [0, 2, 0, 2] layer = normalization.Normalization(axis=-1) layer.adapt(data) output = layer(0.) self.assertListEqual(output.shape.as_list(), [1, 1]) self.assertAllClose(output.numpy(), [[-1]])
def test_broadcasting_during_direct_setting_with_tensors(self): if not tf.executing_eagerly(): self.skipTest("Only supported in TF2.") layer = normalization.Normalization(axis=-1, mean=tf.constant([1.0]), variance=tf.constant([1.0])) output = layer(np.array([[1., 2.]])) expected_output = [[0., 1.]] self.assertAllClose(output, expected_output) self.assertAllClose(layer.get_weights(), [])
def test_mixing_preprocessing_and_regular_layers(self): stage = preprocessing_stage.PreprocessingStage([ image_preprocessing.CenterCrop(16, 16), normalization.Normalization(), convolutional.Conv2D(4, 3) ]) data = np.ones((16, 20, 20, 3), dtype='float32') stage.adapt(data) _ = stage(data) stage.compile('rmsprop', 'mse') stage.fit(data, np.ones((16, 14, 14, 4))) _ = stage.evaluate(data, np.ones((16, 14, 14, 4))) _ = stage.predict(data)
def test_merge_state(self): data = np.random.rand(30, 10, 2) ds = tf.data.Dataset.from_tensor_slices(data).batch(2) norm = normalization.Normalization(axis=(1, 2)) norm.adapt(ds) partial_ds_1 = ds.shard(3, 0) partial_ds_2 = ds.shard(3, 1) partial_ds_3 = ds.shard(3, 2) norm_1 = normalization.Normalization(axis=(1, 2)) norm_2 = normalization.Normalization(axis=(1, 2)) norm_3 = normalization.Normalization(axis=(1, 2)) norm_1.adapt(partial_ds_1) norm_2.adapt(partial_ds_2) norm_3.adapt(partial_ds_3) norm_1.merge_state([norm_2, norm_3]) merged_norm = norm_1 self.assertAllClose(norm(data), merged_norm(data))
def test_1d_unbatched_adapt(self): ds = tf.data.Dataset.from_tensor_slices([ [2.0, 0.0, 2.0, 0.0], [0.0, 2.0, 0.0, 2.0], ]) layer = normalization.Normalization(axis=-1) layer.adapt(ds) output_ds = ds.map(layer) self.assertAllClose( list(output_ds.as_numpy_iterator()), [ [1.0, -1.0, 1.0, -1.0], [-1.0, 1.0, -1.0, 1.0], ], )
def test_layer_computation(self, strategy, adapt_data, axis, test_data, use_dataset, expected): input_shape = tuple([None for _ in range(test_data.ndim - 1)]) if use_dataset: # Keras APIs expect batched datasets adapt_data = tf.data.Dataset.from_tensor_slices(adapt_data).batch( 2) test_data = tf.data.Dataset.from_tensor_slices(test_data).batch(2) with strategy.scope(): input_data = keras.Input(shape=input_shape) layer = normalization.Normalization(axis=axis) layer.adapt(adapt_data) output = layer(input_data) model = keras.Model(input_data, output) output_data = model.predict(test_data) self.assertAllClose(expected, output_data)
def test_mixing_preprocessing_and_regular_layers(self): x0 = Input(shape=(10, 10, 3)) x1 = Input(shape=(10, 10, 3)) x2 = Input(shape=(10, 10, 3)) y0 = merge.Add()([x0, x1]) y1 = image_preprocessing.CenterCrop(8, 8)(x2) y1 = convolutional.ZeroPadding2D(padding=1)(y1) z = merge.Add()([y0, y1]) z = normalization.Normalization()(z) z = convolutional.Conv2D(4, 3)(z) stage = preprocessing_stage.FunctionalPreprocessingStage([x0, x1, x2], z) data = [ np.ones((12, 10, 10, 3), dtype='float32'), np.ones((12, 10, 10, 3), dtype='float32'), np.ones((12, 10, 10, 3), dtype='float32') ] stage.adapt(data) _ = stage(data) stage.compile('rmsprop', 'mse') with self.assertRaisesRegex(ValueError, 'Preprocessing stage'): stage.fit(data, np.ones((12, 8, 8, 4))) ds_x0 = tf.data.Dataset.from_tensor_slices(np.ones((12, 10, 10, 3))) ds_x1 = tf.data.Dataset.from_tensor_slices(np.ones((12, 10, 10, 3))) ds_x2 = tf.data.Dataset.from_tensor_slices(np.ones((12, 10, 10, 3))) ds_x = tf.data.Dataset.zip((ds_x0, ds_x1, ds_x2)) ds_y = tf.data.Dataset.from_tensor_slices(np.ones((12, 8, 8, 4))) dataset = tf.data.Dataset.zip((ds_x, ds_y)).batch(4) with self.assertRaisesRegex(ValueError, 'Preprocessing stage'): stage.fit(dataset) _ = stage.evaluate(data, np.ones((12, 8, 8, 4))) _ = stage.predict(data)
def test_multiple_adapts(self): first_adapt = [[0], [2], [0], [2]] second_adapt = [[2], [4], [2], [4]] predict_input = [[2], [2]] expected_first_output = [[1], [1]] expected_second_output = [[-1], [-1]] inputs = keras.Input(shape=(1,), dtype=tf.int32) layer = normalization.Normalization(axis=-1) layer.adapt(first_adapt) outputs = layer(inputs) model = keras.Model(inputs=inputs, outputs=outputs) actual_output = model.predict(predict_input) self.assertAllClose(actual_output, expected_first_output) # Re-adapt the layer on new inputs. layer.adapt(second_adapt) # Re-compile the model. model.compile() # `predict` should now use the new model state. actual_output = model.predict(predict_input) self.assertAllClose(actual_output, expected_second_output)
def _create_normalization_layer_without_adapt(): return normalization.Normalization(mean=np.random.normal(size=(4, )), variance=np.random.uniform(0.5, 2., size=(4, )))
def _create_normalization_layer_with_adapt(): layer = normalization.Normalization() layer.adapt(np.random.normal(size=(10, 4))) return layer
def test_0d_data(self): layer = normalization.Normalization(axis=None, mean=1.0, variance=1.0) output = layer(0.) self.assertListEqual(output.shape.as_list(), []) self.assertAllClose(output, -1)
def test_keeping_an_unknown_axis_fails(self): layer = normalization.Normalization(axis=-1) with self.assertRaisesRegex(ValueError, "axis.*must have known shape"): layer.build([None])
def test_scalar_input(self): with self.assertRaisesRegex(ValueError, "axis.*values must be in the range"): normalization.Normalization()(1)
def test_broadcasting_during_direct_setting_with_variables_fails(self): with self.assertRaisesRegex(ValueError, "passing a Variable"): _ = normalization.Normalization(axis=-1, mean=tf.Variable([1.0]), variance=tf.Variable([2.0]))
def test_bad_axis_fail_build(self, axis): layer = normalization.Normalization(axis=axis) with self.assertRaisesRegex(ValueError, r"in the range"): layer.build([None, 2, 3])
def test_zeros_fail_init(self, axis): with self.assertRaisesRegex(ValueError, "The argument 'axis' may not be 0."): normalization.Normalization(axis=axis)