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
"""Contains all functions to construct affinity graph for CNC and siamese nets."""
17
from __future__ import absolute_import
18
from __future__ import division
19
from __future__ import print_function
20

21
import numpy as np
22
import tensorflow.compat.v1 as tf
23
from tensorflow.compat.v1.keras import backend as K
24

25

26
def squared_distance(input_x, input_y=None, weight=None):
27
  """Calculates the pairwise distance between points in X and Y.
28

29
  Args:
30
    input_x: n x d matrix
31
    input_y: m x d matrix
32
    weight: affinity n x m -- if provided, we normalize the distance
33

34
  Returns:
35
    n x m matrix of all pairwise squared Euclidean distances
36
  """
37
  if input_y is None:
38
    input_y = input_x
39
  sum_dimensions = list(range(2, K.ndim(input_x) + 1))
40
  input_x = K.expand_dims(input_x, axis=1)
41
  if weight is not None:
42
    # if weight provided, we normalize input_x and input_y by weight
43
    d_diag = K.expand_dims(K.sqrt(K.sum(weight, axis=1)), axis=1)
44
    input_x /= d_diag
45
    input_y /= d_diag
46
  squared_difference = K.square(input_x - input_y)
47
  distance = K.sum(squared_difference, axis=sum_dimensions)
48
  return distance
49

50

51
def knn_affinity(input_x,
52
                 n_nbrs,
53
                 scale=None,
54
                 scale_nbr=None,
55
                 local_scale=None,
56
                 verbose=False):
57
  """Calculates Gaussian affinity matrix.
58

59
  Calculates the symmetrized Gaussian affinity matrix with k1 nonzero
60
  affinities for each point, scaled by
61
  1) a provided scale,
62
  2) the median distance of the k2-th neighbor of each point in X, or
63
  3) a covariance matrix S where S_ii is the distance of the k2-th
64
  neighbor of each point i, and S_ij = 0 for all i != j
65
  Here, k1 = n_nbrs, k2 = scale_nbr
66

67
  Args:
68
    input_x: input dataset of size n
69
    n_nbrs: k1
70
    scale: provided scale
71
    scale_nbr: k2, used if scale not provided
72
    local_scale: if True, then we use the aforementioned option 3), else we
73
      use option 2)
74
    verbose: extra printouts
75

76
  Returns:
77
    n x n affinity matrix
78
  """
79
  if isinstance(n_nbrs, float):
80
    n_nbrs = int(n_nbrs)
81
  elif isinstance(n_nbrs,
82
                  tf.Variable) and n_nbrs.dtype.as_numpy_dtype != np.int32:
83
    n_nbrs = tf.cast(n_nbrs, np.int32)
84
  # get squared distance
85
  dist_x = squared_distance(input_x)
86
  # calculate the top k losest neighbors
87
  nn = tf.nn.top_k(-dist_x, n_nbrs, sorted=True)
88

89
  vals = nn[0]
90
  # apply scale
91
  if scale is None:
92
    # if scale not provided, use local scale
93
    if scale_nbr is None:
94
      scale_nbr = 0
95
    else:
96
      assert scale_nbr > 0 and scale_nbr <= n_nbrs
97
    if local_scale:
98
      scale = -nn[0][:, scale_nbr - 1]
99
      scale = tf.reshape(scale, [-1, 1])
100
      scale = tf.tile(scale, [1, n_nbrs])
101
      scale = tf.reshape(scale, [-1, 1])
102
      vals = tf.reshape(vals, [-1, 1])
103
      if verbose:
104
        vals = tf.Print(vals, [tf.shape(vals), tf.shape(scale)],
105
                        'vals, scale shape')
106
      vals = vals / (2 * scale)
107
      vals = tf.reshape(vals, [-1, n_nbrs])
108
    else:
109

110
      def get_median(scales, m):
111
        with tf.device('/cpu:0'):
112
          scales = tf.nn.top_k(scales, m)[0]
113
        scale = scales[m - 1]
114
        return scale, scales
115

116
      scales = -vals[:, scale_nbr - 1]
117
      const = tf.shape(input_x)[0] // 2
118
      scale, scales = get_median(scales, const)
119
      vals = vals / (2 * scale)
120
  else:
121
    # otherwise, use provided value for global scale
122
    vals = vals / (2 * scale**2)
123

124
  # get the affinity
125
  aff_vals = tf.exp(vals)
126
  # flatten this into a single vector of values to shove in a sparse matrix
127
  aff_vals = tf.reshape(aff_vals, [-1])
128
  # get the matrix of indices corresponding to each rank
129
  # with 1 in the first column and k in the kth column
130
  nn_ind = nn[1]
131
  # get the j index for the sparse matrix
132
  j_index = tf.reshape(nn_ind, [-1, 1])
133
  # the i index is just sequential to the j matrix
134
  i_index = tf.range(tf.shape(nn_ind)[0])
135
  i_index = tf.reshape(i_index, [-1, 1])
136
  i_index = tf.tile(i_index, [1, tf.shape(nn_ind)[1]])
137
  i_index = tf.reshape(i_index, [-1, 1])
138
  # concatenate the indices to build the sparse matrix
139
  indices = tf.concat((i_index, j_index), axis=1)
140
  # assemble the sparse weight matrix
141
  weight_mat = tf.SparseTensor(
142
      indices=tf.cast(indices, dtype='int64'),
143
      values=aff_vals,
144
      dense_shape=tf.cast(tf.shape(dist_x), dtype='int64'))
145
  # fix the ordering of the indices
146
  weight_mat = tf.sparse_reorder(weight_mat)
147
  # convert to dense tensor
148
  weight_mat = tf.sparse_tensor_to_dense(weight_mat)
149
  # symmetrize
150
  weight_mat = (weight_mat + tf.transpose(weight_mat)) / 2.0
151

152
  return weight_mat
153

154

155
def full_affinity(input_x, scale):
156
  """Calculates the symmetrized full Gaussian affinity matrix, scaled by a provided scale.
157

158
  Args:
159
    input_x: input dataset of size n x d
160
    scale: provided scale
161

162
  Returns:
163
    n x n affinity matrix
164
  """
165
  sigma = K.variable(scale)
166
  dist_x = squared_distance(input_x)
167
  sigma_squared = K.expand_dims(K.pow(sigma, 2), -1)
168
  weight_mat = K.exp(-dist_x / (2 * sigma_squared))
169
  return weight_mat
170

171

172
def get_contrastive_loss(m_neg=1, m_pos=.2):
173
  """Contrastive loss from Hadsell-et-al.'06.
174

175
  http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf.
176

177
  Args:
178
    m_neg: negativeness.
179
    m_pos: possitiveness.
180

181
  Returns:
182
    Contrastive loss
183
  """
184

185
  def contrastive_loss(y_true, y_pred):
186
    return K.mean(y_true * K.square(K.maximum(y_pred - m_pos, 0)) +
187
                  (1 - y_true) * K.square(K.maximum(m_neg - y_pred, 0)))
188

189
  return contrastive_loss
190

191

192
def euclidean_distance(vects):
193
  """Computes the euclidean distances between vects[0] and vects[1]."""
194
  x, y = vects
195
  return K.sqrt(
196
      K.maximum(K.sum(K.square(x - y), axis=1, keepdims=True), K.epsilon()))
197

198

199
def eucl_dist_output_shape(shapes):
200
  """Provides the output shape of the above computation."""
201
  s_1, _ = shapes
202
  return (s_1[0], 1)
203