def test_layer_invocation(self, use_gate, num_blocks, dropout_position,
dtype):
tf.keras.mixed_precision.experimental.set_policy(dtype)
kwargs = dict(intermediate_size=16,
intermediate_activation="relu",
dropout=0.1,
use_gate=use_gate,
num_blocks=num_blocks,
dropout_position=dropout_position,
kernel_initializer="glorot_uniform",
bias_initializer="zeros")
test_layer = gated_feedforward.GatedFeedforward(**kwargs)
sequence_length = 16
width = 32
# Create a 3-dimensional input (the first dimension is implicit).
data_tensor = tf.keras.Input(shape=(sequence_length, width))
output_tensor = test_layer(data_tensor)
# Create a model from the test layer.
model = tf.keras.Model(data_tensor, output_tensor)
# Invoke the model on test data.
batch_size = 6
input_data = 10 * np.random.random_sample(
(batch_size, sequence_length, width))
output_data = model.predict(input_data)
self.assertEqual(output_data.shape,
(batch_size, sequence_length, width))
def test_serialize_deserialize(self):
kwargs = dict(intermediate_size=16,
intermediate_activation="relu",
dropout=0.1,
use_gate=False,
num_blocks=4,
dropout_position="after_residual",
kernel_initializer="glorot_uniform",
bias_initializer="zeros")
test_layer = gated_feedforward.GatedFeedforward(**kwargs)
new_layer = gated_feedforward.GatedFeedforward.from_config(
test_layer.get_config())
# If the serialization was successful, the new config should match the old.
self.assertAllEqual(test_layer.get_config(), new_layer.get_config())
def test_layer_creation(self, use_gate, num_blocks, dropout_position, dtype):
tf.keras.mixed_precision.set_global_policy(dtype)
kwargs = dict(
inner_dim=128,
inner_activation="relu",
dropout=0.1,
use_gate=use_gate,
num_blocks=num_blocks,
dropout_position=dropout_position,
kernel_initializer="glorot_uniform",
bias_initializer="zeros")
test_layer = gated_feedforward.GatedFeedforward(**kwargs)
sequence_length = 64
width = 128
# Create a 3-dimensional input (the first dimension is implicit).
data_tensor = tf.keras.Input(shape=(sequence_length, width))
output_tensor = test_layer(data_tensor)
# The default output of a transformer layer should be the same as the input.
self.assertEqual(data_tensor.shape.as_list(), output_tensor.shape.as_list())