python - 如何保存从冻结推理图生成的 TensorRT 图？

Question

我使用下面的脚本将我的 freeze_inference_graph 转换为 TensorRT 优化的：

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

with tf.Session() as sess:
    # First deserialize your frozen graph:
    with tf.gfile.GFile('frozen_inference_graph.pb', 'rb') as f:
        frozen_graph = tf.GraphDef()
        frozen_graph.ParseFromString(f.read())
    # Now you can create a TensorRT inference graph from your
    # frozen graph:
    converter = trt.TrtGraphConverter(
        input_graph_def=frozen_graph,
        nodes_blacklist=['outputs/Softmax']) #output nodes
    trt_graph = converter.convert()
    # Import the TensorRT graph into a new graph and run:
    output_node = tf.import_graph_def(
        trt_graph,
        return_elements=['outputs/Softmax'])
    sess.run(output_node)

我的问题是如何将这个优化的图形保存到磁盘上以便我可以使用它来运行推理？

Answer 1

是的，您只需添加这两行：

saved_model_dir_trt = "./tensorrt_model.trt"
converter.save(saved_model_dir_trt)