knn-classifier memory leak (peak bytes) #8159
Description
System information
- I have written custom code
- Linux Ubuntu 22.04
- TensorFlow.js installed from npm
- TensorFlow.js version: 4.16.0
- knn-classifier version: 1.2.6
- Node version: 20.11.0
Describe the current behavior
Running predictClass for first time with "many" example classes consumes all memory and kills node process.
Describe the expected behavior
Running predictClass for first time with "many" example classes will consume fewer peak bytes during processing, allowing scripts to load more classes (depends on specific usage).
Standalone code to reproduce the issue
https://stackblitz.com/edit/stackblitz-starters-gnuzh5?file=index.js
Other info / logs Include any logs or source code that would be helpful to
Run locally I am seeing peak bytes grow quickly. With 16GiB RAM it cannot complete with 1000 classes.
Proposed Fix
I think the issue is in
https://github.com/tensorflow/tfjs-models/blob/3a8b906f5a95f8a2d697c3896f20b7529cdeaf6f/knn-classifier/src/index.ts#L106-L109
for (const label in this.classDatasetMatrices) {
newTrainLogitsMatrix = concatWithNulls(
newTrainLogitsMatrix, this.classDatasetMatrices[label]);
}Each concatWithNulls allocates a new tensor with the previous concatenation remaining in memory. Although this is called within a tidy so the memory will be released after the similarities have been computed, during computation a large footprint is required.
Instead, the tensors could be disposed of in the loop.
for (const label in this.classDatasetMatrices) {
const newTrainLogitsMatrixToBeDisposed = newTrainLogitsMatrix;
newTrainLogitsMatrix = concatWithNulls(
newTrainLogitsMatrix, this.classDatasetMatrices[label]);
newTrainLogitsMatrixToBeDisposed?.dispose();
}With that change the peak bytes look much better.
Although, better performance could probably be found by replacing the concat operation entirely....