如何将Tensorflow .pb模型转换为Tensforflow Lite

时间:2020-03-19 15:12:24

标签: tensorflow model

我需要使用Google CoLab tensorflow pb 模型转换为 tensorflow lite

下一个转换步骤:

1)上传模型:

from google.colab import files
pbfile = files.upload()

2)进行转换:

import tensorflow as tf
pb_file = 'data_513.pb'
tflite_file = 'data_513.tlite'

converter = tf.lite.TFLiteConverter.from_frozen_graph(pb_file, ['ImageTensor'], ['SemanticPredictions'], 
                                                      input_shapes={"ImageTensor":[1,513,513,3]})

tflite_model = converter.convert()
open(tflite_file,'wb').write(tflite_model) 

转换失败并出现下一个错误

检查失败:array.data_type == array.final_data_type阵列“ ImageTensor”的实际和最终数据类型(data_type = uint8,final_data_type = float)不匹配。

我认为我可能需要指定一些额外的命令来克服此错误,但是我找不到关于它的任何信息。

1 个答案:

答案 0 :(得分:1)

最后找到了解决方案。此处已被其他人使用:

import tensorflow as tf
pb_file = 'model.pb'
tflite_file = 'model.tflite'

converter = tf.lite.TFLiteConverter.from_frozen_graph(pb_file, ['ImageTensor'], ['SemanticPredictions'], 
                                                      input_shapes={"ImageTensor":[1,513,513,3]})


converter.inference_input_type=tf.uint8
converter.quantized_input_stats = {'ImageTensor': (128, 127)}  # (mean, stddev)

tflite_model = converter.convert()
open(tflite_file,'wb').write(tflite_model) 

interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

files.download(tflite_file)