google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""Tests for supcon.classification_head."""
17
from absl.testing import parameterized
18

19
import numpy as np
20
import tensorflow.compat.v1 as tf
21

22
from supcon import classification_head
23

24

25
class ClassificationHeadTest(tf.test.TestCase, parameterized.TestCase):
26

27
  @parameterized.named_parameters(
28
      ('rank_1', 1),
29
      ('rank_4', 4),
30
      ('rank_8', 8),
31
  )
32
  def testIncorrectRank(self, rank):
33
    inputs = tf.compat.v1.placeholder(tf.float32, shape=[10] * rank)
34
    with self.assertRaisesRegex(ValueError, 'is expected to have rank 2'):
35
      classifier = classification_head.ClassificationHead(num_classes=10)
36
      classifier(inputs)
37

38
  @parameterized.named_parameters(
39
      ('float32', tf.float32),
40
      ('float64', tf.float64),
41
      ('float16', tf.float16),
42
  )
43
  def testConstructClassificationHead(self, dtype):
44
    batch_size = 3
45
    num_classes = 10
46
    input_shape = [batch_size, 4]
47
    expected_output_shape = [batch_size, num_classes]
48
    inputs = tf.random.uniform(input_shape, seed=1, dtype=dtype)
49
    classifier = classification_head.ClassificationHead(num_classes=num_classes)
50
    output = classifier(inputs)
51
    self.assertListEqual(expected_output_shape, output.shape.as_list())
52
    self.assertEqual(inputs.dtype, output.dtype)
53

54
  def testGradient(self):
55
    inputs = tf.random.uniform((3, 4), dtype=tf.float64, seed=1)
56
    classifier = classification_head.ClassificationHead(num_classes=10)
57
    output = classifier(inputs)
58
    gradient = tf.gradients(output, inputs)
59
    self.assertIsNotNone(gradient)
60

61
  def testCreateVariables(self):
62
    inputs = tf.random.uniform((3, 4), dtype=tf.float64, seed=1)
63
    classifier = classification_head.ClassificationHead(num_classes=10)
64
    classifier(inputs)
65
    self.assertLen(
66
        [var for var in tf.trainable_variables() if 'kernel' in var.name], 1)
67
    self.assertLen(
68
        [var for var in tf.trainable_variables() if 'bias' in var.name], 1)
69

70
  def testInputOutput(self):
71
    batch_size = 3
72
    num_classes = 10
73
    expected_output_shape = (batch_size, num_classes)
74
    inputs = tf.random.uniform((batch_size, 4), dtype=tf.float64, seed=1)
75
    classifier = classification_head.ClassificationHead(num_classes=num_classes)
76
    output_tensor = classifier(inputs)
77
    with self.cached_session() as sess:
78
      sess.run(tf.compat.v1.global_variables_initializer())
79
      outputs = sess.run(output_tensor)
80
      # Make sure that there are no NaNs
81
      self.assertFalse(np.isnan(outputs).any())
82
      self.assertEqual(outputs.shape, expected_output_shape)
83

84

85
if __name__ == '__main__':
86
  tf.test.main()
87