Coral Edge TPU编译器无法转换tflite模型:模型未量化

时间:2019-07-03 12:49:54

标签: python tensorflow keras tensorflow-lite tpu

我正在尝试使用TensorFlow lite部署一个简单的测试应用程序。我想在我的设备上使用Coral Edge TPU Stick,所以我必须执行量化意识培训。我想适合一个功能f(x) = 2 x - 1。我的训练代码如下:

import tensorflow as tf
import numpy as np
from tensorflow import keras
from tensorflow.contrib import lite

# Create model
model = keras.models.Sequential([keras.layers.Dense(units=1, input_shape=[1])])

# Quantization aware training
sess = keras.backend.get_session()
tf.contrib.quantize.create_training_graph(sess.graph)
sess.run(tf.global_variables_initializer())

tf.summary.FileWriter('logs/', graph=sess.graph)

model.compile(optimizer='sgd', loss='mean_squared_error')

# Training data
xs = np.array([ -1.0, 0.0, 1.0, 2.0, 3.0, 4.0], dtype=float)
ys = np.array([ -3.0, -1.0, 1.0, 3.0, 5.0, 7.0], dtype=float)

model.fit(xs, ys, epochs=500, batch_size=2)

# Test the model for plausbility
print(model.predict([10.0]))

# Display the quantization-relevant variables
for node in sess.graph.as_graph_def().node:
    if 'weights_quant/AssignMaxLast' in node.name \
        or 'weights_quant/AssignMinLast' in node.name:
        tensor = sess.graph.get_tensor_by_name(node.name + ':0')
        print('{} = {}'.format(node.name, sess.run(tensor)))


# Save the keras model
keras_file = 'quant_linear.h5'
keras.models.save_model(model, keras_file)

# Convert the keras model into a tflite model
converter = lite.TocoConverter.from_keras_model_file(keras_file)
converter.post_training_quantize = True
tflite_model = converter.convert()
open('quant_linear.tflite', 'wb').write(tflite_model)

作为输出,我得到了(省略了keras和CUDA的特定输出):

[[18.86733]]
dense/weights_quant/AssignMinLast = 0.0
dense/weights_quant/AssignMaxLast = 1.984399676322937

这里需要注意的两件事:该模型是合理的,它应该输出接近19的值。显然,它也使用量化权重。如果我未启用量化意识训练,则这两个变量将不会显示。

此外,此模型可以由tf-lite解释器实例加载和执行。为了能够在TPU支持下使用它,我必须使用tpuedge_compiler对其进行转换。安装后,我执行

edgetpu_compiler quant_linear.tflite

不幸的是,似乎无法识别该模型已量化。输出

user@ubuntu:~/TensorFlow$ edgetpu_compiler quant_linear.tflite 
Edge TPU Compiler version 1.0.249710469
INFO: Initialized TensorFlow Lite runtime.
Invalid model: quant_linear.tflite
Model not quantized

我试图在线编译它,但也失败了。这是错误还是在训练/转换过程中我搞砸了?另外,也许有工具可以验证我是否确实使用了量化模型?

谢谢!

1 个答案:

答案 0 :(得分:0)

在TFLite转换AFAIK期间必须使用显式quantization。 量化Keras模型的代码示例:

dataset = tf.data.Dataset(...)

def generator():
    for item in dataset:
        image = # get image from dataset item
        yield [np.array([image.astype(np.float32)])]

converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8
converter.representative_dataset = tf.lite.RepresentativeDataset(generator)

model = converter.convert()