API解析
常用的就这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'])