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.

Experiment setup

The experiment has 2 main classes: TfKNN and TfliteKNN . 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 TfKNN
• Run nearest neighbors search using TfKNN
• Export tflite model from TfKNN
• Load tflite model and run nearest neighbors search using TfliteKNN
• Compare nearest neighbors search results generated from TfKNN and TfliteKNN

Background configurations

• The experiment uses tensorflow 2.4.0
• There are a couple of constants used for illustration:

K = 3: the number of neighbors found for each KNN search
N_FEATURES = 2: the number of features for input data
N_SAMPLES = 1000: the number of samples (dataset size)
N_CENTERS = 5: the number of clusters drawn from the synthetic data
RANDOM_STATE = 0: seed value for deterministic experiment 
TEST_SIZE = 0.3: ratio of data size split for test set

Tensorflow KNN

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:

• TfKNN needs 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 TfKNN is l2 distance.
• TfKNN.neighbors is 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).

Evaluation

Dataset

The dataset is generated by using make_blobs() function from sklearn.datasets . The dataset size is 1000, and the test ratio is 0.3.

Dataset generate by `make_blobs()`

Evaluation process

• Step 1: training data is enrolled into TfKNN
• Step 2: tflite model is exported from TfKNN
• Step 3: run knn search on both TfKNN and TfliteKNN
• Step 4: compare search results on test data from both implementations

The above snippet ran the experiment and confirmed the results from tf_knnand 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 :)