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 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

## 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.

## 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_knn`

and `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 :)