TensorFlow Lite开发系列之python接口解析(一)

环境: tensorflow2.x， 一定要使用linux系统，后期转换模型windows会出现bug

API解析

网址：https://tensorflow.google.cn/api_docs/python/tf/lite

API	功能
class Interpreter	推理TensorFlow Lite模型
class TFLiteConverter	将TensorFlow模型转换为TensorFlow Lite模型
class Optimize	定义tflite图使用此API的优化
class OpsSet	TFLite模型的操作集

常用的就这4个类，常用方法介绍：

# model_path：TF-Lite Flatbuffer文件的路径
# model_content：模型的内容
tf.lite.Interpreter(
    model_path=None, model_content=None, experimental_delegates=None
)

allocate_tensors()   # 加载所有的tensor

get_input_details()  # 获取模型输入详细信息

get_output_details()  # 获取模型输出详细信息

 # 获取输入张量的值（获取副本），该值可以从get_output_details中的“索引”字段中获得
get_tensor( 
    tensor_index
)

get_tensor_details()   # 返回值：包含张量信息的字典列表

invoke()  # 进行推理， 在调用它之前，请确保设置输入大小，分配张量和填充值

# tensor_index：要设置的张量的张量索引。该值可以从get_input_details中的“索引”字段中获得
# value：要设置的张量值
set_tensor(
    tensor_index, value
)

# Converting a SavedModel to a TensorFlow Lite model.
converter = lite.TFLiteConverter.from_saved_model(saved_model_dir)
tflite_model = converter.convert()

# Converting a tf.Keras model to a TensorFlow Lite model.
converter = lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

# Converting ConcreteFunctions to a TensorFlow Lite model.
converter = lite.TFLiteConverter.from_concrete_functions([func])
tflite_model = converter.convert()

tf.lite.Optimize 
类变量
	DEFAULT
	OPTIMIZE_FOR_LATENCY
	OPTIMIZE_FOR_SIZE

完整例子：

1. 训练一个模型用于后面转换TFLite模型与推理

import tensorflow as tf
from tensorflow import keras

# 读取数据
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = tf.expand_dims(x_train, 3)
y_train = keras.utils.to_categorical(y_train, num_classes=10)
datasets = tf.data.Dataset.from_tensor_slices((x_train, y_train))
datasets = datasets.repeat(1).batch(10)

# 定义模型
img = keras.Input(shape=[28, 28, 1])

x = keras.layers.Conv2D(filters=64, kernel_size=4, strides=1, padding='SAME', activation='relu')(img)
x = keras.layers.AveragePooling2D(pool_size=2, strides=2, padding='SAME')(x)
x = keras.layers.BatchNormalization()(x)
x = keras.layers.Dropout(0.15)(x)

x = keras.layers.Flatten()(x)
x = keras.layers.Dense(512, activation='relu')(x)
y_pred = keras.layers.Dense(10, activation='softmax')(x)

model = keras.Model(inputs=img, outputs=y_pred)

model.compile(optimizer=keras.optimizers.Adam(0.01),
              loss=keras.losses.categorical_crossentropy,
              metrics=['AUC', 'accuracy'])

model.fit(datasets, epochs=1)

model.save('./model.h5')

2. 转换模型

import tensorflow as tf
from tensorflow import keras

model = tf.keras.models.load_model("./model.h5")
model.summary()

converter = tf.lite.TFLiteConverter.from_keras_model(model)  # 生成转化器
tflite_model = converter.convert()  # 进行转换
open('./converted_model.tflite', 'wb').write(tflite_model)  # 写入

3. 推理

import tensorflow as tf
from tensorflow import keras
import numpy as np

interpreter = tf.lite.Interpreter(model_path='./converted_model.tflite')  # 读入并生成interpreter
interpreter.allocate_tensors()  # 加载所有的张量

input_details = interpreter.get_input_details()  # 获取输入的信息
output_details = interpreter.get_output_details()  # 获取输出的信息

# 指定随机数进行预测
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()

# 根据tensor索引取出推理的结果
tflite_result = interpreter.get_tensor(output_details[0]['index'])