Tensorflow Keras模型的量化意识训练

时间:2020-03-31 01:25:09

标签: python tensorflow keras quantization-aware-training

我想用我的keras模型进行量化训练。我已经尝试过如下。我正在使用tensorflow 1.14.0

train_graph = tf.Graph()
train_sess = tf.compat.v1.Session(graph=train_graph)
tf.compat.v1.keras.backend.set_session(train_sess)

with train_graph.as_default():
    tf.keras.backend.set_learning_phase(1)
    model = my_keras_model()

    tf.contrib.quantize.create_training_graph(input_graph = train_graph, quant_delay=5)
    train_sess.run(tf.global_variables_initializer()) 

    model.compile(...)
    model.fit_generator(...)

    saver = tf.compat.v1.train.Saver()  
    saver.save(train_sess, checkpoint_path)

它可以正常工作。

但是,保存的模型(h5和ckpt)的大小与没有量化的模型完全相同。

这是正确的方法吗?如何检查量化是否正确?

或者,有更好的量化方法吗?

1 个答案:

答案 0 :(得分:1)

完成量化感知训练并将模型保存到磁盘时,实际上尚未对其进行量化。换句话说,它是为量化而“准备”的,但是权重仍然是float32。您必须将模型进一步转换为TFLite才能真正进行量化。您可以使用以下代码来实现:

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]

quantized_tflite_model = converter.convert()

这将通过int8权重和uint8激活来量化您的模型。

请查看official example,以供进一步参考。