Skip to content

Commit

Permalink
Embedding Projector: reduce knn computation (#6321)
Browse files Browse the repository at this point in the history
## Motivation for features / changes
for t-SNE max number of knn neighbors is 3 * max_perplexity (3 * 100 =
300). For umap max number of neighbors is 100.
let's just compute 300 knn neighbors so we can slice the result each
time when we tweak perplexity for tSNE.

## Technical description of changes
When computing KNN, use an upper threshold of 300 so any future
computation with same number of points will be cached.


## Screenshots of UI changes

N/A

## Detailed steps to verify changes work correctly (as executed by you)
1. Launch local app
2. Open t-SNE
3. adjust to Perplexity > 25
4. verify knn isn't computed a second time

## Alternate designs / implementations considered
  • Loading branch information
alicialics authored Apr 19, 2023
1 parent d1c504f commit 4f0cd50
Showing 1 changed file with 15 additions and 3 deletions.
18 changes: 15 additions & 3 deletions tensorboard/plugins/projector/vz_projector/data.ts
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,8 @@ export const PCA_SAMPLE_DIM = 200;
const NUM_PCA_COMPONENTS = 10;
/** Id of message box used for umap optimization progress bar. */
const UMAP_MSG_ID = 'umap-optimization';
/** Minimum KNN neighbors threshold */
const MIN_NUM_KNN_NEIGHBORS = 300;
/**
* Reserved metadata attributes used for sequence information
* NOTE: Use "__seq_next__" as "__next__" is deprecated.
Expand Down Expand Up @@ -474,16 +476,26 @@ export class DataSet {
);
} else {
const knnGpuEnabled = (await util.hasWebGLSupport()) && !IS_FIREFOX;
const numKnnNeighborsToCompute = Math.max(
nNeighbors,
MIN_NUM_KNN_NEIGHBORS
);
const result = await (knnGpuEnabled
? knn.findKNNGPUCosDistNorm(data, nNeighbors, (d) => d.vector)
? knn.findKNNGPUCosDistNorm(
data,
numKnnNeighborsToCompute,
(d) => d.vector
)
: knn.findKNN(
data,
nNeighbors,
numKnnNeighborsToCompute,
(d) => d.vector,
(a, b) => vector.cosDistNorm(a, b)
));
this.nearest = result;
return Promise.resolve(result);
return Promise.resolve(
result.map((neighbors) => neighbors.slice(0, nNeighbors))
);
}
}
/* Perturb TSNE and update dataset point coordinates. */
Expand Down

0 comments on commit 4f0cd50

Please sign in to comment.