Skip to content

Commit

Permalink
Embedding Projector: perform consistent knn calculation
Browse files Browse the repository at this point in the history
  • Loading branch information
alicialics committed Apr 20, 2023
1 parent 4f0cd50 commit 18cf7c5
Showing 1 changed file with 6 additions and 14 deletions.
20 changes: 6 additions & 14 deletions tensorboard/plugins/projector/vz_projector/data.ts
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,7 @@ export interface DataPoint {
[key: string]: number;
};
}
const IS_FIREFOX = navigator.userAgent.toLowerCase().indexOf('firefox') >= 0;

/** Controls whether nearest neighbors computation is done on the GPU or CPU. */
export const TSNE_SAMPLE_SIZE = 10000;
export const UMAP_SAMPLE_SIZE = 5000;
Expand Down Expand Up @@ -475,23 +475,15 @@ export class DataSet {
.map((neighbors) => neighbors.slice(0, nNeighbors))
);
} else {
const knnGpuEnabled = (await util.hasWebGLSupport()) && !IS_FIREFOX;
const numKnnNeighborsToCompute = Math.max(
nNeighbors,
MIN_NUM_KNN_NEIGHBORS
);
const result = await (knnGpuEnabled
? knn.findKNNGPUCosDistNorm(
data,
numKnnNeighborsToCompute,
(d) => d.vector
)
: knn.findKNN(
data,
numKnnNeighborsToCompute,
(d) => d.vector,
(a, b) => vector.cosDistNorm(a, b)
));
const result = await knn.findKNNGPUCosDistNorm(
data,
numKnnNeighborsToCompute,
(d) => d.vector
);
this.nearest = result;
return Promise.resolve(
result.map((neighbors) => neighbors.slice(0, nNeighbors))
Expand Down

0 comments on commit 18cf7c5

Please sign in to comment.