def test_layer_invocation(self):
test_layer = cls_head.ClassificationHead(inner_dim=5, num_classes=2)
features = tf.zeros(shape=(2, 10, 10), dtype=tf.float32)
output = test_layer(features)
self.assertAllClose(output, [[0., 0.], [0., 0.]])
self.assertSameElements(test_layer.checkpoint_items.keys(),
["pooler_dense"])
def test_pooler_layer(self, inner_dim, num_weights_expected):
test_layer = cls_head.ClassificationHead(inner_dim=inner_dim, num_classes=2)
features = tf.zeros(shape=(2, 10, 10), dtype=tf.float32)
_ = test_layer(features)
num_weights_observed = len(test_layer.get_weights())
self.assertEqual(num_weights_observed, num_weights_expected)
def test_layer_serialization(self):
layer = cls_head.ClassificationHead(10, 2)
new_layer = cls_head.ClassificationHead.from_config(layer.get_config())
# If the serialization was successful, the new config should match the old.
self.assertAllEqual(layer.get_config(), new_layer.get_config())