我编写了一个基于VGG16的模型,只添加了两个额外的卷积层。输出是一个大小为16x16x1的数组,仅是简单二进制分类的结果。我使用了TensorFlow-lite,并根据可用文档编写了代码。问题在于,当我使用该模型进行预测时,它需要很长时间(近5分钟)才能给出结果。
我在GPU上使用Tensorflow 2.4,Python 3.7,我的显卡是GTX 1660Ti(移动版),CPU是intel i7 9750H。
以下是可用的代码。
import tensorflow as tf
import os
import time
import numpy as np
import cv2
import keras
import pathlib
saved_model_dir= 'model/'
saved_modelh5 = 'model.h5'
dataset_path = 'bound box dataset/img'
out_path = 'converted_model.tflite'
num_calibration_steps = 10
#-----------------------------------------------------------
images = []
for file in os.listdir(dataset_path):
img = cv2.imread( os.path.join(dataset_path,file) )
images.append(img)
images = np.array( images )
imgs_tensor = tf.cast( images, dtype = tf.float32)/255.0
ds = tf.data.Dataset.from_tensor_slices((imgs_tensor)).batch(1)
print('data loaded')
#-----------------------------------------------------------
def representative_dataset_gen():
for input_value in ds.take(num_calibration_steps):
yield [input_value]
#converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter = tf.lite.TFLiteConverter.from_keras_model(keras.models.load_model(saved_modelh5))
converter.optimizations = [tf.lite.Optimize.DEFAULT ]
#converter.representative_dataset = tf.lite.RepresentativeDataset( representative_dataset_gen )
#converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
tflite_model = converter.convert()
#------------------------------------------------------------
#with open(out_path, "wb")as f:
# f.write(tflite_model)
print('converted')
tflite_model_file = pathlib.Path(out_path)
tflite_model_file.write_bytes(tflite_model)
print('Saved')
img = cv2.imread('bound box dataset/img/1.png')
input_data = img.reshape(1,512,512,3).astype(np.float32)/255.0
interpreter = tf.lite.Interpreter( model_content = tflite_model)
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.
t = time.time()
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'])
t = time.time() - t
print('predict time:',t)
converter.optimizations
或者converter.representative_dataset
时,预测时间会上升到大约5分钟左右。但是当我不初始化它们时,预测时间降至大约1秒钟左右。你有什么想法为什么会这样吗? - ARK1375