如何在脚本中加载tflite模型？

Question

如何在脚本中加载tflite模型？

pythontensorflowtensorflow-lite

26

我已经使用 bazel 将 .pb 文件转换为 tflite 文件。现在我想在我的Python脚本中加载这个 tflite 模型，只是为了测试是否会给我正确的输出。

- Harshit Mishra

2个回答

5

在Python中使用TensorFlow Lite模型：

TensorFlow Lite的冗长性很强大，因为它允许更多的控制，但在许多情况下，您只想传递输入并获得输出，因此我创建了一个包装这个逻辑的类：

以下内容适用于来自tfhub.dev的分类模型，例如：https://tfhub.dev/tensorflow/lite-model/mobilenet_v2_1.0_224/1/metadata/1

# Usage
model = TensorflowLiteClassificationModel("path/to/model.tflite")
(label, probability) = model.run_from_filepath("path/to/image.jpeg")

import tensorflow as tf
import numpy as np
from PIL import Image


class TensorflowLiteClassificationModel:
    def __init__(self, model_path, labels, image_size=224):
        self.interpreter = tf.lite.Interpreter(model_path=model_path)
        self.interpreter.allocate_tensors()
        self._input_details = self.interpreter.get_input_details()
        self._output_details = self.interpreter.get_output_details()
        self.labels = labels
        self.image_size=image_size

    def run_from_filepath(self, image_path):
        input_data_type = self._input_details[0]["dtype"]
        image = np.array(Image.open(image_path).resize((self.image_size, self.image_size)), dtype=input_data_type)
        if input_data_type == np.float32:
            image = image / 255.

        if image.shape == (1, 224, 224):
            image = np.stack(image*3, axis=0)

        return self.run(image)

    def run(self, image):
        """
        args:
          image: a (1, image_size, image_size, 3) np.array

        Returns list of [Label, Probability], of type List<str, float>
        """

        self.interpreter.set_tensor(self._input_details[0]["index"], image)
        self.interpreter.invoke()
        tflite_interpreter_output = self.interpreter.get_tensor(self._output_details[0]["index"])
        probabilities = np.array(tflite_interpreter_output[0])

        # create list of ["label", probability], ordered descending probability
        label_to_probabilities = []
        for i, probability in enumerate(probabilities):
            label_to_probabilities.append([self.labels[i], float(probability)])
        return sorted(label_to_probabilities, key=lambda element: element[1])

注意

但是，您需要修改此内容以支持不同的用例，因为我正在传递图像作为输入，并获得分类（[标签，概率]）输出。如果您需要文本输入（NLP），或其他输出（物体检测输出边界框、标签和概率），分类（仅标签）等。

此外，如果您期望不同大小的图像输入，则需要更改输入大小并重新分配模型(self.interpreter.allocate_tensors())。这很慢（低效）。最好使用平台调整大小功能（例如Android图形库）来进行调整大小，而不是使用TensorFlow lite模型进行调整大小。或者，您可以使用一个单独的模型来调整模型的大小，这样可以更快地allocate_tensors()。

- Ben Butterworth

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Jing Zhao · Accepted Answer

你可以使用 TensorFlow Lite Python 解释器在 Python shell 中加载 tflite 模型，并使用输入数据测试它。

代码将会像这样：

import numpy as np
import tensorflow as tf

# Load TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model_path="converted_model.tflite")
interpreter.allocate_tensors()

# Get input and output tensors.
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# Test model on random input data.
input_shape = input_details[0]['shape']
input_data = np.array(np.random.random_sample(input_shape), dtype=np.float32)
interpreter.set_tensor(input_details[0]['index'], input_data)

interpreter.invoke()

# The function `get_tensor()` returns a copy of the tensor data.
# Use `tensor()` in order to get a pointer to the tensor.
output_data = interpreter.get_tensor(output_details[0]['index'])
print(output_data)

以上代码来自 TensorFlow Lite 官方指南，如需更详细的信息，请阅读此处。