Tensorflow量化:数组输出不具有MinMax信息

时间:2019-05-14 21:11:46

标签: python tensorflow tensorflow-lite

我正在尝试创建一个Tensorflow量化模型,以使用Coral USB Accelerator进行推理。这是我的问题的一个最小的独立示例:

import sys

import tensorflow as tf

CKPT = "a/out.ckpt"
TFLITE = "a/out.tflite"

args = sys.argv[1:]
if 0 == len(args):
    print("Options are 'train' or 'save'")
    exit(-1)
cmd = args[0]

if cmd not in ["train", "save"]:
    print("Options are 'train' or 'save'")
    exit(-1)

tr_in = [[1.0, 0.0], [0.0, 1.0], [0.0, 0.0], [1.0, 1.0]]
tr_out = [[1.0], [1.0], [0.0], [0.0]]

nn_in = tf.placeholder(tf.float32, (None, 2), name="input")

W = tf.Variable(tf.random_normal([2, 1], stddev=0.1))
B = tf.Variable(tf.ones([1]))

nn_out = tf.nn.relu6(tf.matmul(nn_in, W) + B, name="output")

if "train" == cmd:
    tf.contrib.quantize.create_training_graph(quant_delay=0)
    nn_act = tf.placeholder(tf.float32, (None, 1), name="actual")
    diff = tf.reduce_mean(tf.pow(nn_act - nn_out, 2))
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
        optimizer = tf.train.AdamOptimizer(
            learning_rate=0.0001,
        )
        goal = optimizer.minimize(diff)
else:
    tf.contrib.quantize.create_eval_graph()

init = tf.global_variables_initializer()
with tf.Session() as session:
    session.run(init)

    saver = tf.train.Saver()
    try:
        saver.restore(session, CKPT)
    except BaseException as e:
        print("While trying to restore: {}".format(str(e)))

    if "train" == cmd:
        for epoch in range(2):
            _, d = session.run([goal, diff], feed_dict={
                nn_in: tr_in,
                nn_act: tr_out,
            })
            print("Loss: {}".format(d))
        saver.save(session, CKPT)
    elif "save" == cmd:
        converter = tf.lite.TFLiteConverter.from_session(
            session, [nn_in], [nn_out],
        )
        converter.inference_type = tf.lite.constants.QUANTIZED_UINT8
        input_arrays = converter.get_input_arrays()
        converter.quantized_input_stats = {input_arrays[0] : (0.0, 1.0)}
        tflite_model = converter.convert()
        with open(TFLITE, "wb") as f:
            f.write(tflite_model)

假设您有一个名为“ a”的目录,则可以使用以下目录运行该文件:

python example.py train
python example.py save

“火车”步骤应该可以正常工作,但是当尝试导出量化的tflite文件时,我得到以下信息:

...
2019-05-14 14:03:44.032912: F tensorflow/lite/toco/graph_transformations/quantize.cc:144] Array output does not have MinMax information, and is not a constant array. Cannot proceed with quantization.
Aborted

我的目标是成功运行“保存”步骤并得到训练有素的量化模型。我想念什么?

1 个答案:

答案 0 :(得分:0)

TFLiteConverter中有一个棘手的错误:

  • 要转换为量化模型格式,每个(几乎每个)数学运算节点都需要其他节点(带有MinMax信息)。
  • 在相应操作之后添加了 create_eval_graph 函数的此类附加节点。
  • 但是在转换为TFLite格式的过程中,转换器仅考虑了输入输出之间的节点(包括两端)。因此,在这种情况下,您的 nn_out 之后的其他节点(带有MinMax信息)将被“扔掉”,从而导致上述转换错误:(

如果您构建的分类网络通常以softmax层结尾(不需要MinMax信息),则不会出现该错误。但是对于回归网络,这是一个问题。我使用以下解决方法。

在调用 create_eval_graph 函数之前,在输出层之后添加其他(实际上无意义的)操作,如下所示:

nn_out = tf.minimum(nn_out, 1e6)

您可以使用任何任意数字(用于第二个参数),该数字要比预期的输出层值上限大得多。就我而言,它非常有效。