python调用tflite模型

最新推荐文章于 2024-07-12 07:47:18 发布

落花逐流水

最新推荐文章于 2024-07-12 07:47:18 发布

阅读量4.8k

点赞数 1

分类专栏： tensorflow python

本文链接：https://blog.csdn.net/juluwangriyue/article/details/118082260

版权

python 同时被 2 个专栏收录

66 篇文章 0 订阅

订阅专栏

tensorflow

9 篇文章 0 订阅

订阅专栏

制作frozen模型

模型制作参考前面的一篇博客《tensorflow 20：搭网络、导出模型、运行模型》。主要就是两层卷积和两层全连接用来识别mnist数据集，保存为frozen模型文件。

转换为tflite模型

非量化转换
转换代码：

# -*- coding:utf-8 -*-
import tensorflow as tf

in_path = "./model/frozen_graph.pb"
out_path = "./model/frozen_graph.tflite"
# out_path = "./model/quantize_frozen_graph.tflite"

# 模型输入节点
input_tensor_name = ["input/x"]
input_tensor_shape = {"input/x":[1, 784]}
# 模型输出节点
classes_tensor_name = ["out/fc2"]

converter = tf.lite.TFLiteConverter.from_frozen_graph(in_path,
                                            input_tensor_name, classes_tensor_name,
                                            input_shapes = input_tensor_shape)
#converter.post_training_quantize = True
tflite_model = converter.convert()

with open(out_path, "wb") as f:
    f.write(tflite_model)

转换模型前后，模型文件大小几乎一样，都是12M左右。

量化转换

把上面代码里‘converter.post_training_quantize = True’启用就行了。
转换出的模型大小变为原来的约1/4, 只有3M左右。

PC上用python调用tflite模型

调用非量化模型

# -*- coding:utf-8 -*-
import os
os.environ["CUDA_VISIBLE_DEVICES"]="-1"  
import cv2
import numpy as np
import time

import tensorflow as tf

test_image_dir = './test_images/'
#model_path = "./model/quantize_frozen_graph.tflite"
model_path = "./model/frozen_graph.tflite"

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

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

#with tf.Session( ) as sess:
if 1:
    file_list = os.listdir(test_image_dir)
    
    model_interpreter_time = 0
    start_time = time.time()
    # 遍历文件
    for file in file_list:
        print('=========================')
        full_path = os.path.join(test_image_dir, file)
        print('full_path:{}'.format(full_path))
        
        # 只要黑白的，大小控制在(28,28)
        img = cv2.imread(full_path, cv2.IMREAD_GRAYSCALE )
        res_img = cv2.resize(img,(28,28),interpolation=cv2.INTER_CUBIC) 
        # 变成长784的一维数据
        new_img = res_img.reshape((784))
        
        # 增加一个维度，变为 [1, 784]
        image_np_expanded = np.expand_dims(new_img, axis=0)
        image_np_expanded = image_np_expanded.astype('float32') # 类型也要满足要求
        
        # 填装数据
        model_interpreter_start_time = time.time()
        interpreter.set_tensor(input_details[0]['index'], image_np_expanded)
        
        # 注意注意，我要调用模型了
        interpreter.invoke()
        output_data = interpreter.get_tensor(output_details[0]['index'])
        model_interpreter_time += time.time() - model_interpreter_start_time
        
        # 出来的结果去掉没用的维度
        result = np.squeeze(output_data)
        print('result:{}'.format(result))
        #print('result:{}'.format(sess.run(output, feed_dict={newInput_X: image_np_expanded})))
        
        # 输出结果是长度为10（对应0-9）的一维数据，最大值的下标就是预测的数字
        print('result:{}'.format( (np.where(result==np.max(result)))[0][0]  ))
    used_time = time.time() - start_time
    print('used_time:{}'.format(used_time))
    print('model_interpreter_time:{}'.format(model_interpreter_time))