How to calculate squared sum quickly in Tensorflow2?

Question

I wanna calculate this with huge input data.

I did it with tensorflow. but it is calculated one by one. so it is not fast enough. I tried to make fixed_mat and input_np like (1000, 101, 1088). it didn't work... is there any suggestions to calculate huge input??? Thanks for advance.

fixed_mat.shape

(101, 1088)

input_np.shape

(1000, 1088)

import tensorflow as tf
num = 2
res = tf.reduce_sum(tf.math.squared_difference(fixed_mat, 
np.array([input_np[1]] * fixed_mat.shape[0])), 1)
vals, indice = tf.nn.top_k(tf.negative(res), num)
print(list(indice.numpy()), list(- vals.numpy()))

[13, 90] [1.3422034332504837, 1.8790145150615656]

Answer 1

If you want to speed up tensorflow code, you should not use the eager mode. It is best to use the @tf.function decorator

See this guide: https://www.tensorflow.org/guide/function

This code works:

@tf.function
def my_tf_function(fixed_mat, input_np, num):
    res = tf.reduce_sum(tf.math.squared_difference(
        fixed_mat,
        tf.gather(input_np, 1)), 1)
    vals, indice = tf.nn.top_k(tf.negative(res), num)
    return vals, indice