使用TensorRT（TF-TRT）将张量流save_model从float32转换为float16的问题

Question

我有一个tensorflow（版本1.14）float32 SavedModel，我想转换为float16。根据https://docs.nvidia.com/deeplearning/frameworks/tf-trt-user-guide/index.html#usage-example，我可以将“ FP16”传递给precision_mode以将模型转换为fp16。但是在检查了张量板之后，转换后的模型仍然是fp32：净参数是DT_FLOAT而不是DT_HALF。转换后的模型的大小类似于转换前的模型。 （在这里，我假设，如果成功转换，由于参数将被切成两半，因此模型将变大一半）。

import tensorflow as tf
from tensorflow.python.compiler.tensorrt import trt_convert as trt
import os

FLAGS = tf.flags.FLAGS
tf.flags.DEFINE_string('saved_model_dir', '', 'Input saved model dir.')
tf.flags.DEFINE_bool('use_float16', False,
                     'Whether we want to quantize it to float16.')
tf.flags.DEFINE_string('output_dir', '', 'Output saved model dir.')


def main(argv):
    del argv  # Unused.
    saved_model_dir = FLAGS.saved_model_dir
    output_dir = FLAGS.output_dir
    use_float16 = FLAGS.use_float16

    precision_mode = "FP16" if use_float16 else "FP32"
    converter = trt.TrtGraphConverter(input_saved_model_dir=saved_model_dir,
                                      precision_mode=precision_mode)
    converter.convert()
    converter.save(output_dir)


if __name__ == '__main__':
    tf.app.run(main)

非常欢迎任何意见或建议！谢谢

Answer 1

请尝试通过更改：

tf.flags.DEFINE_bool('use_float16', False, 'Whether we want to quantize it to float16.')

到

tf.flags.DEFINE_bool('use_float16', True, 'Whether we want to quantize it to float16.')

这应该起作用或给出适当的错误日志，因为使用当前代码precision_mode被设置为"FP32"。您需要precision_mode = "FP16"才能测试半精度。