CODEWITHSUNDEEP
pythontensorflow
magmacollaris
magmacollaris

Reputation: 89

In Tensorflow, how can I create a new tensor from a cosine similarity operation?

I'm trying to create an output tensor with dimensionality 32 x 576 x 2 from an operation between matrices M and X, with the following shapes:

M.shape: (576, 2, 2048)
X.shape: (32, 2048)

The operation I'm defining is an element-wise cosine similarity, from the following equation:

enter image description here

which represents the cosine similarity between the feature vector 𝑥 and the vector M_j,k.

This is how I've implemented it in code (incorrectly), where BATCH_SIZE=32, C=576, V=2:

    @tf.function
    def call(self, X):
        M = self.kernel

        norm_M = tf.norm(M, ord=2, axis=2)
        norm_X = tf.norm(X, ord=2, axis=1)

        l_r = (some scalar value, separate to this question)

        # Compute cosine similarity between X and M
        # as a matrix with dimensionality:
        # BATCH_SIZE x C x V
        feature_batch_size = tf.shape(X)[0]
        c = tf.shape(M)[0]
        v = tf.shape(M)[1]
        output_matrix = tf.zeros([feature_batch_size, c, v])
        output_matrix = tf.Variable(output_matrix, trainable=False)
        for row in tf.range(feature_batch_size):
          for column in tf.range(c):
            for channel in tf.range(v):
              a = tf.tensordot(M[column][channel], X[row], 1)
              b = norm_M[column][channel] * norm_X[row]
              output_matrix[row][column][channel] = a / b

        return [output_matrix, l_r]

This fails on the line output_matrix[row][column][channel] = a / b because it's unhappy with an assignment to an individual row:column:channel of a tf.Variable.

Is there a better way to do this operation over these two matrices to create the desired output matrix so that it can be done without these three nested for loops and maintain compatibility with the tf.Function graph functionality?

If not, what can I do to assign variables to individual elements on a tf.Variable as I'm unsuccessfully attempting to do here?

Extra information:
norm_M.shape: (576, 2)
norm_X.shape: (32,)

Upvotes: 0

Views: 95

Answers (1)

Vijay Mariappan
Vijay Mariappan

Reputation: 17201

You can replace these loops completely by using vectorized operations in the place of for loops.

num = tf.einsum('ij,klj->ikl',X,M)
denom = tf.einsum('i,jk->ijk',norm_X, norm_M)
output_matrix = num/denom

Upvotes: 2

Related Questions