javascripttensorflowtensorflowjs

TensorflowJS: Optimal way to calculate distance or similarity between multiple tensors?


I'm writing an algorithm that requires I compare two different arrays of tensors, dataset and centroids. dataset has +1000 more tensors than centroids does, and all tensors have the same dimension ([1 x n]).

My current solution (code below) is as follows: For each tensor in dataset, find the distance between that tensor and all tensors in centroids, then store the index of the closest centroid.

dataset.forEach(d => {
      const distances = centroids.map(centroid => getEuclidianDistance(d.tensor, centroid));
      const centroidIndex = distances.indexOf(Math.min(...distances));
      d.centroid = centroidIndex;
    })

This works, but is pretty slow. It's also a nested loop, which feels inefficient.

Is there a better way to do this with tensorflowjs (i.e. via some sort of similarity matrix?).

Thanks!

P.S. - If a particular solution requires a particular distance function, I'm totally open to changing mine. Currently my distance function is the following:

    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];
    }

Solution

  • I had a similar requirement a few months back in which, given a 2d point, I was looking to find from an array of line segments which was closest to the point. I struggled trying to coerce tensorflowjs to perform this efficiently, and eventually stumbled upon gpu.js which is more geared towards compiling custom GPU kernel functions.

    In the example below that I cooked up, I have arrays representing 11 (X,Y) coordinates and another pair of arrays representing 5 (X,Y) coordinates. The result will be a matrix 11x5 which computes every distance between both sets of points. The key function is "kernel", which is compiled by gpu.js to operate on a GPU core, and essentially computes the distance between a pair of points sourced from both the 11 coordinates and 5 coordinates. In theory, this kernel function will be laid across many GPU cores to accelerate performance. Ie, in this case perform all 55 at once. (I say "in theory", because gpu.js as I understand it leverages the webGL shader map function, and am not entirely sure of the impact of the virtualization layers involved in the stack that leads to the GPU cores actually performing the work...)

    The result is an 11x5 matrix containing the distance from each combination of point pairs, whereupon this 11x5 matrix is then piped to "kernelMin", which will be a bit slower, as it is looping through the results seeking the minimum value, and also capturing the index of the minimum value. That being said, there should be 11 concurrent GPU cores working away on the effort of finding which of the 5 coordinates is closest.

    const kernel = gpu.createKernel(function(x0, y0, x1, y1) {
      let dx = x1[this.thread.y][0] - x0[0][this.thread.x];
      let dy = y1[this.thread.y][0] - y0[0][this.thread.x];
      return Math.sqrt(dx * dx + dy * dy);
    }).setPipeline(true).setOutput([11,5]);
    
    const result1 = kernel(
      GPU.input(
        new Float32Array([0,10,20,30,40,50,60,70,80,90,100]),
        [11,1]
      ),
      GPU.input(
        new Float32Array([100,100,100,100,100,100,100,100,100,100,100]),
        [11,1]
      ),
      GPU.input(
        new Float32Array([0,30,50,70,100]),
        [1,5]
      ),
      GPU.input(
        new Float32Array([10,10,10,10,10]),
        [1,5]
      )
    );
    
    console.log(result1.toArray());
    
    const kernelMin = gpu.createKernel(function(a) {
      let minVal = 1000000;
      let minIdx = 0;
      for (let y = 0; y < 5; y++) {
        if (a[y][this.thread.x] < minVal) {
          minVal = a[y][this.thread.x];
          minIdx = y;
        }
      }
      return [minVal,minIdx];
    }).setOutput([11]);
    
    const result2 = kernelMin(result1);
    
    console.log(result2);
    

    The final output is...

    0: Float32Array(2) [90, 0]
    1: Float32Array(2) [90.55384826660156, 0]
    2: Float32Array(2) [90.55384826660156, 1]
    3: Float32Array(2) [90, 1]
    4: Float32Array(2) [90.55384826660156, 1]
    5: Float32Array(2) [90, 2]
    6: Float32Array(2) [90.55384826660156, 2]
    7: Float32Array(2) [90, 3]
    8: Float32Array(2) [90.55384826660156, 3]
    9: Float32Array(2) [90.55384826660156, 4]
    10: Float32Array(2) [90, 4]
    

    Note that I've hard coded the matrix sizes into the example for clarity. Gpu.js obviously accepts variables. Also, in your case, depending on the size of your matrices, you might have to break the problem into chunks depending on the amount of GPU RAM required to house the complete cross matrix of distances...

    I realize this isn't tensorflowjs, but hope this helps.

    EDIT - via TensorFlow.JS

    Spent a few minutes porting to tensorflow.js. The core concept is to build the matrices of x and y values in preparation for performing calculations en mass.

    const x0 = tf.tensor1d([0,10,20,30,40,50,60,70,80,90,100]);
    const y0 = tf.tensor1d([100,100,100,100,100,100,100,100,100,100,100]);
    
    const x1 = tf.tensor1d([0,30,50,70,100]);
    const y1 = tf.tensor1d([10,10,10,10,10]);
    
    const x0mat = x0.tile([5]).reshape([5,11]);
    const y0mat = y0.tile([5]).reshape([5,11]);
    const x1mat = x1.tile([11]).reshape([11,5]).transpose();
    const y1mat = y1.tile([11]).reshape([11,5]).transpose();
    
    x0mat.print();
    x1mat.print();
    const xDeltas = x1mat.squaredDifference(x0mat);
    
    y0mat.print();
    y1mat.print();
    const yDeltas = y1mat.squaredDifference(y0mat);
    
    const distance = xDeltas.add(yDeltas).sqrt();
    distance.print();
    
    distance.argMin(1).print();
    distance.min(1).print();
    

    With results of...

    Tensor - x0mat
        [[0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100],
         [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100],
         [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100],
         [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100],
         [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100]]
    Tensor - x1mat
        [[0  , 0  , 0  , 0  , 0  , 0  , 0  , 0  , 0  , 0  , 0  ],
         [30 , 30 , 30 , 30 , 30 , 30 , 30 , 30 , 30 , 30 , 30 ],
         [50 , 50 , 50 , 50 , 50 , 50 , 50 , 50 , 50 , 50 , 50 ],
         [70 , 70 , 70 , 70 , 70 , 70 , 70 , 70 , 70 , 70 , 70 ],
         [100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100]]
    Tensor - y0mat
        [[100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100],
         [100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100],
         [100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100],
         [100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100],
         [100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100]]
    Tensor - y1mat
        [[10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10],
         [10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10],
         [10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10],
         [10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10],
         [10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10]]
    Tensor - distance
        [[90         , 90.5538483 , 92.1954422 , 94.8683319, 98.4885788 , 102.9562988, 108.1665344, 114.01754 , 120.415947 , 127.2792206, 134.5362396],
         [94.8683319 , 92.1954422 , 90.5538483 , 90        , 90.5538483 , 92.1954422 , 94.8683319 , 98.4885788, 102.9562988, 108.1665344, 114.01754  ],
         [102.9562988, 98.4885788 , 94.8683319 , 92.1954422, 90.5538483 , 90         , 90.5538483 , 92.1954422, 94.8683319 , 98.4885788 , 102.9562988],
         [114.01754  , 108.1665344, 102.9562988, 98.4885788, 94.8683319 , 92.1954422 , 90.5538483 , 90        , 90.5538483 , 92.1954422 , 94.8683319 ],
         [134.5362396, 127.2792206, 120.415947 , 114.01754 , 108.1665344, 102.9562988, 98.4885788 , 94.8683319, 92.1954422 , 90.5538483 , 90         ]]
    Tensor - argMin of distance
        [0, 3, 5, 7, 10]
    Tensor - min of distance
        [90, 90, 90, 90, 90]
    

    The code is broken up in steps to show the basic concept. I'm sure it can be condensed and optimized further.