Using TF backend, I need to construct a similarity matrices of two 3D vectors, both with shape (batch_size, N, M), being N and M natural numbers.
The function tf.losses.cosine_distance is only between 1D tensors. I need to build a Tensor matrix batch_sizexNxM such that matrix[k][i][j] will be the cosine similarity of the Tensor1[k][i] and Tensor2[k][j].
How can I do this?
I know of no pairwise distance operations in Keras or tensorflow. But the matrix math can be implemented in TF/Keras backend code and then placed in a Keras layer.
Here's the matrix representation of the cosine similarity of two vectors: $$ cos(\theta) = \frac{\mathbf{A}\cdot\mathbf{B}}{\|\mathbf{A}\|_2 \|\mathbf{B}\|_2} $$
I'll show the code and a test that confirms that it works. First, generate non-trivial test data.
import numpy as np
import keras
import keras.backend as K
# set up test data
n_batch = 100
n = 400 # number of points in the first set
m = 500 # number of points in the second set
d = 200 # number of dimensions
A = np.random.rand(n_batch, n, d)
B = np.random.rand(n_batch, m, d)
Define pairwise cosine similarity function.
# convenience l2_norm function
def l2_norm(x, axis=None):
"""
takes an input tensor and returns the l2 norm along specified axis
"""
square_sum = K.sum(K.square(x), axis=axis, keepdims=True)
norm = K.sqrt(K.maximum(square_sum, K.epsilon()))
return norm
def pairwise_cosine_sim(A_B):
"""
A [batch x n x d] tensor of n rows with d dimensions
B [batch x m x d] tensor of n rows with d dimensions
returns:
D [batch x n x m] tensor of cosine similarity scores between each point i<n, j<m
"""
A, B = A_B
A_mag = l2_norm(A, axis=2)
B_mag = l2_norm(B, axis=2)
num = K.batch_dot(A_tensor, K.permute_dimensions(B_tensor, (0,2,1)))
den = (A_mag * K.permute_dimensions(B_mag, (0,2,1)))
dist_mat = num / den
return dist_mat
Build a shallow Keras model around the function.
# build dummy model
A_tensor = K.constant(A)
B_tensor = K.constant(B)
A_input = keras.Input(tensor=A_tensor)
B_input = keras.Input(tensor=B_tensor)
dist_output = keras.layers.Lambda(pairwise_cosine_sim)([A_input, B_input])
dist_model = keras.Model(inputs=[A_input, B_input], outputs=dist_output)
dist_model.compile("sgd", "mse")
Compare to sklearn implementation
sk_dist = np.zeros( (n_batch, n, m) )
for i in range(n_batch):
sk_dist[i,...] = cosine_similarity(A[i,...], B[i,...])
keras_dist = dist_model.predict(None, steps=1)
np.allclose(sk_dist, keras_dist)
> True
Lambda don't accept two arguments. You could solve this by making your pairwise_cosine receive the arguments in a list instead of separated. However there is another issue. I need this layer to accept 3D Tensors actually, where the 1st dimension is the batch size. If they were 2D Tensors I could indeed solve it as you said. But in this case I would need something like Lambda(lambda x, y: pairwise_cosine(x, y)). Sorry for not mentioning this previously, was trying to solve one problem at a time.
$\endgroup$
A_mag and B_mag afterwards. Can you refactor your answer so that A and B are of the form (?, n, m), where ? is the batch size?
$\endgroup$
pairwise_cosine_sim function so that it receive [? x n x d] and [? x m x d]
$\endgroup$