Running K-Nearest Neighbors on Tensorflow Lite
Recently, I have started working more towards running Machine Learning (ML) models on edge devices. Tensorflow Lite is a great library for running Deep Learning (DL) models cross platform. However, for classical ML models, I found it was a little difficult to find a solution for it at first.
After diving a little deeper into how Tensorflow Lite operates and how to convert from Tensorflow to Tensorflow Lite, I did some experiments for running classical ML models, especially at inference. I found it was possible to write mathematical operations using Tensorflow and later convert to Tensorflow Lite.
This article uses K-Nearest Neighbors as an example, but the same idea can apply for running other algorithms.
The experiment has 2 main classes:
TfKNN (Tensorflow KNN) is responsible for the actual Tensorflow implementation of KNN and
TfliteKNN (Tensorflow Lite) is used to load the exported
tflite file from
TfKNN and run prediction on it.
The experiment illustrates the following points:
- Enroll training data to
- Run nearest neighbors search using
tflitemodel and run nearest neighbors search using
- Compare nearest neighbors search results generated from
- The experiment uses
- There are a couple of constants used for illustration:
K = 3: the number of neighbors found for each KNN searchN_FEATURES = 2: the number of features for input dataN_SAMPLES = 1000: the number of samples (dataset size)N_CENTERS = 5: the number of clusters drawn from the synthetic dataRANDOM_STATE = 0: seed value for deterministic experiment TEST_SIZE = 0.3: ratio of data size split for test set
Since KNN is a lazy learning algorithm, the inference (search process) requires access to the enrolled data (training data). There are a couple of points that worth mentioning:
TfKNNneeds to take in the training data (
train_tensor) as an attribute in order to run the search operation at inference.
- The distance function used in
TfKNNis l2 distance.
TfKNN.neighborsis the actual function that performs KNN search. Also, after TF lite conversion, this is the method executed by the tflite model.
Tensorflow Lite KNN
TfliteKNN is a class which encapsulates the loading and running of tflite Interpreter.
TfliteKNN.neighbors is functionally equivalent to
TfKNN.neighbors (perform nearest neighbors search).
The dataset is generated by using
make_blobs() function from
sklearn.datasets . The dataset size is 1000, and the test ratio is 0.3.
- Step 1: training data is enrolled into
- Step 2: tflite model is exported from
- Step 3: run knn search on both
- Step 4: compare search results on test data from both implementations
The above snippet ran the experiment and confirmed the results from
tflite_knn are identical.
Extra step — use KNN for clustering
Besides the comparison evaluation, one additional experiment that was done was to use KNN for clustering. Specifically, since
make_blobs() provides an option to specify the number of clusters to generate the data, we can utilize that for doing clustering.
The results of the above experiment are:
Tf KNN accuracy:0.9433333333333334
TfLite KNN accuracy: 0.9433333333333334
The accuracy isn’t too bad, is it?
Happy coding :)