Save TensorFlowJS MobileNet + KNN to TFLite

I have trained a KNN on top of MobileNet logits results using TensorFlowJS.

And I want to know how can I export the result of the MobileNet + KNN to a TFLite model.

const knn = knnClassifier.create()
const net = await mobilenet.load()

const handleTrain = (imgEl, label) => {
  const image = tf.browser.fromPixels(imgEl);
  const activation = net.infer(image, true);
  knn.addExample(activation, label)
}

Answer

1. Save the model

Save the model this example saves the file to the native file system or if you need it to be saved in other places then check the documentation.

await model.save('file:///path/to/my-model');

You should have a JSON file and a binary weight file(s) after this step.

2. Convert from TensorFlow.js Layers model to Saved Model format

tfjs_model.json is the path to the model.json that you get from the previous step and saved_model is the path where you want to save the SavedModel format.
You can read more about using the TensorflowJS Converter from here.

tensorflowjs_converter --input_format=tfjs_layers_model --output_format=keras_saved_model tfjs_model.json saved_model

3. Convert from SavedModel format to TFLite format

Converting from a SavedModel format to TFLite is the recommended way to do this as per the documentation.

import tensorflow as tf

# Convert the model
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir) # path to the SavedModel directory
tflite_model = converter.convert()

# Save the model.
with open('model.tflite', 'wb') as f:
  f.write(tflite_model)