TensorflowJS：计算多个张量之间的距离或相似性的最佳方法？

Question

我正在编写一个算法，需要我比较两个不同的张量数组，dataset并且centroids. dataset有 +1000 个以上的张量centroids，并且所有张量都具有相同的维度 ( [1 x n])。

我当前的解决方案（下面的代码）如下：对于 中的每个张量dataset，找到该张量与 中所有张量之间的距离centroids，然后存储最近的 的索引centroid。

dataset.forEach(d => {

      const distances = centroids.map(centroid => getEuclidianDistance(d.tensor, centroid));

      const centroidIndex = distances.indexOf(Math.min(...distances));

      d.centroid = centroidIndex;

    })

这有效，但速度很慢。它也是一个嵌套循环，感觉效率低下。

有没有更好的方法来使用 tensorflowjs（即通过某种相似性矩阵？）。

谢谢！

PS - 如果一个特定的解决方案需要一个特定的距离函数，我完全愿意改变我的。目前我的距离函数如下：

    getEuclidianDistance(arr1, arr2) {

        // calculate euclidian distance between two arrays

        let distTensor = tf.tidy(() => {

            const distance = tf.squaredDifference(arr1, arr2).sum().sqrt();

            return distance.dataSync()

        })

        return distTensor[0];

    }