在 python 中加载 Tensorflow Lite 模型

Posted 2023-02-15

【中文标题】在 python 中加载 Tensorflow Lite 模型

【英文标题】：Load Tensorflow Lite models in python

【发布时间】：2021-09-08 22:23:26

【问题描述】：

我正在使用具有量化和浮点模型的 Tensorflow Lite 进行 TinyML 项目。在我的管道中，我使用tf.keras API 训练我的模型，然后将模型转换为 TFLite 模型。最后，我将 TFLite 模型量化为 int8。 我可以使用 API model.save 和 tf.keras.model.load_model 保存和加载“正常”张量流模型

是否可以对转换后的 TFLite 模型执行相同的操作？每次都经过量化过程非常耗时。

【问题讨论】：

你的意思是你想在转换后直接从 tflite 模型中获取推理，而不是在其他设备上部署它？

是的，我需要一种将 TFLite 模型保存在磁盘上的方法，以便我可以在第二天加载它们并在我的 python 笔记本中进行推理。目前，我只能保存和加载 tensorflow 模型，但我总是必须通过量化才能使用 TFLite 进行推理

【参考方案1】：

您可以使用tflite interpreter 直接在笔记本中从 TFLite 模型中获取推理。

这是图像分类模型的示例。假设我们有一个 tflite 模型：

tflite_model_file = 'converted_model.tflite'

然后我们可以像这样加载和测试它：

# Load TFLite model and allocate tensors.
with open(tflite_model_file, 'rb') as fid:
    tflite_model = fid.read()
    
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

input_index = interpreter.get_input_details()[0]["index"]
output_index = interpreter.get_output_details()[0]["index"]

# Gather results for the randomly sampled test images
predictions = []

test_labels, test_imgs = [], []
for img, label in tqdm(test_batches.take(10)):
    interpreter.set_tensor(input_index, img)
    interpreter.invoke()
    predictions.append(interpreter.get_tensor(output_index))
    
    test_labels.append(label.numpy()[0])
    test_imgs.append(img)

请注意，您只能从 tflite 模型进行推断。您无法更改架构和层，例如重新加载 Keras 模型。如果你想改变架构，你应该保存 Keras 模型，并对其进行测试，直到你得到满意的结果，然后将其转换为 tflite。

【讨论】：

我没有意识到我可以将 tflite 模型加载为二进制文件。有时，简单的解决方案非常简单，以至于很难发现它们。谢谢！