Loading TFlite model for Inference (Python)

Question

I'm using Tensorflow Lite to train an image classifier. I now have a bunch of *.tflite models stored, and I'm trying to write some code that allows me to pick a tflite model file, pick a dataset, and test that model on that dataset (inference).

When I train a model using:

model = image_classifier.create(trainData, validation_data=valData, shuffle=True, use_augmentation=False)

I am able to easily test this model on a test dataset right away because the model is actually stored in the variable 'model', by using:

model.evaluate_tflite('model.tflite', test_data)

or

loss, accuracy = model.evaluate(test_data)

However, if I simply want to load an already existing *.tflite model, without having trained it in the same run, I can't figure out a simple way to do that.

Following these instructions, it seems to be a lot of steps for what I'm trying to do. In other Machine Learning libraries (like PyTorch), you are able to define the model and then quickly load the saved weights and then get to testing, like:

model = models.densenet201(progress=True, pretrained=pretrained)
model.load_state_dict(torch.load("models/model.pt"))

Is there a simple way for me to initialise the model into the 'model' variable, load the saved weights from a *.tflite file, and then run inference?

Thank you for your help

u1234x1234 · Answer

A simple example of image classification:

import tensorflow as tf
import numpy as np
import cv2


class TFLiteModel:
    def __init__(self, model_path: str):
        self.interpreter = tf.lite.Interpreter(model_path)
        self.interpreter.allocate_tensors()

        self.input_details = self.interpreter.get_input_details()
        self.output_details = self.interpreter.get_output_details()

    def predict(self, *data_args):
        assert len(data_args) == len(self.input_details)
        for data, details in zip(data_args, self.input_details):
            self.interpreter.set_tensor(details["index"], data)
        self.interpreter.invoke()
        return self.interpreter.get_tensor(self.output_details[0]["index"])


model = TFLiteModel("mobilenet_v2_1.0_224_1_default_1.tflite")
image = cv2.imread("hand_blower.png")
image = cv2.resize(image, (224, 224))
image = image.astype(np.float32)[np.newaxis]
image = (image - 127.5) / 127.5

label = model.predict(image)[0].argmax()
print(label)

Please refer to the official documentation for detailed information:

https://www.tensorflow.org/api_docs/python/tf/lite/Interpreter

The model was loaded from:

https://tfhub.dev/tensorflow/lite-model/mobilenet_v2_1.0_224/1/default/1

Loading TFlite model for Inference (Python)

1 Answers