def get_keras_model(num_classes, input_length, use_batchnorm=True, l2=1e-5, num_clusters=None, alpha_init=None): """Make a model.""" assert not num_clusters or not alpha_init model = tf.keras.models.Sequential() model.add(tf.keras.Input((input_length, ))) # Input is [bs, input_length] trill_layer = hub.KerasLayer( handle= 'https://tfhub.dev/google/nonsemantic-speech-benchmark/trill-distilled/3', trainable=True, arguments={'sample_rate': tf.constant(16000, tf.int32)}, output_key='embedding', output_shape=[None, 2048]) assert trill_layer.trainable_variables model.add(trill_layer) if num_clusters and num_clusters > 0: model.add(NetVLAD(num_clusters=num_clusters)) if use_batchnorm: model.add(tf.keras.layers.BatchNormalization()) elif alpha_init is not None: model.add( autopool.AutoPool(axis=1, alpha_init=alpha_init, trainable=False)) else: model.add(tf.keras.layers.Lambda(lambda x: tf.reduce_mean(x, axis=1))) model.add( tf.keras.layers.Dense( num_classes, kernel_regularizer=tf.keras.regularizers.l2(l=l2))) return model
def test_keras_model_has_no_trainable_vars(self): """Test that pooling variable is properly trainable.""" m = tf.keras.models.Sequential( [tf.keras.Input(shape=(4, 5)), autopool.AutoPool(axis=1, trainable=False)] ) m.build() self.assertEmpty(m.trainable_variables)
def test_limiting_behavior(self, alpha_init, agg_func): """Test semantics of alpha. When `alpha` = 0, it reduces to an unweighted mean; when `alpha` = 1, it simplifies to soft-max pooling; and when `alpha` -> inf, it approaches the max operator, and `alpha` -> -inf, it approachs the min operator. Args: alpha_init: The alpha parameters. agg_func: The expected aggregation function. """ a = tf.random.uniform(shape=(2, 3, 1), seed=12) actual = autopool.AutoPool(axis=1, alpha_init=alpha_init)(a, keepdims=True) expected = agg_func(a, axis=1, keepdims=True) self.assertAllClose(actual, expected)
def get_keras_model(num_classes, use_batchnorm=True, l2=1e-5, num_clusters=None, alpha_init=None): """Make a model.""" model = tf.keras.models.Sequential() if num_clusters and num_clusters > 0: model.add(NetVLAD(num_clusters=num_clusters)) if use_batchnorm: model.add(tf.keras.layers.BatchNormalization()) elif alpha_init is not None: model.add(autopool.AutoPool(axis=1, alpha_init=alpha_init, trainable=False)) else: model.add(tf.keras.layers.Lambda(lambda x: tf.reduce_mean(x, axis=1))) model.add(tf.keras.layers.Dense( num_classes, kernel_regularizer=tf.keras.regularizers.l2(l=l2))) return model
def test_keras_model_has_trainable_vars(self): """Test that pooling variable is properly trainable.""" m = tf.keras.models.Sequential( [tf.keras.Input(shape=(4, 5)), autopool.AutoPool(axis=1)]) m.build() self.assertIsNotNone(m.trainable_variables) assert len(m.trainable_variables) == 1 # Run a fake training step and check that values after the step are # different. Access through `average_alpha` to check that it works. original_var_val = m.get_layer(index=0).average_alpha.numpy() with tf.GradientTape() as tape: o = m(tf.random.uniform(shape=(3, 4, 5)), training=True) o.shape.assert_is_compatible_with((3, 5)) loss_value = tf.keras.losses.CategoricalCrossentropy( from_logits=True)(y_true=tf.ones_like(o), y_pred=o) grads = tape.gradient(loss_value, m.trainable_variables) tf.keras.optimizers.Adam(learning_rate=1.0).apply_gradients( zip(grads, m.trainable_variables)) post_training_var_val = m.get_layer(index=0).average_alpha.numpy() self.assertNotAllClose(original_var_val, post_training_var_val)
def test_average_alpha(self): a = tf.random.uniform(shape=(3, 4, 5), seed=123) layer = autopool.AutoPool(axis=1) layer(a) # Triggers layer build. self.assertIsNotNone(layer.average_alpha) self.assertEqual(layer.average_alpha.shape, [])