对于使用 TensorRT 优化的 TensorFlow 图,我遇到了极长的加载时间。未优化的加载速度很快,但加载优化的加载相同的代码需要10 多分钟:
trt_graph_def = tf.GraphDef()
with tf.gfile.GFile(pb_path, 'rb') as pf:
trt_graph_def.ParseFromString(pf.read())
我在 NVIDIA Drive PX 2 设备上(如果有的话),TensorFlow 1.12.0 是从源代码、CUDA 9.2 和 TensorRT 4.1.1 构建的。由于它卡在 ParseFromString() 我怀疑 protobuf 所以这里是它的配置:
$ dpkg -l | grep protobuf
ii libmirprotobuf3:arm64 0.26.3+16.04.20170605-0ubuntu1.1 arm64 Display server for Ubuntu - RPC definitions
ii libprotobuf-dev:arm64 2.6.1-1.3 arm64 protocol buffers C++ library (development files)
ii libprotobuf-lite9v5:arm64 2.6.1-1.3 arm64 protocol buffers C++ library (lite version)
ii libprotobuf9v5:arm64 2.6.1-1.3 arm64 protocol buffers C++ library
ii protobuf-compiler 2.6.1-1.3 arm64 compiler for protocol buffer definition files
$ pip3 freeze | grep protobuf
protobuf==3.6.1
这是我将非优化模型转换为 TRT 模型的方式:
def get_frozen_graph(graph_file):
"""Read Frozen Graph file from disk."""
with tf.gfile.FastGFile(graph_file, "rb") as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
return graph_def
print("Load frozen graph from disk")
frozen_graph = get_frozen_graph(DATA_DIR + MODEL + '.pb')
print("Optimize the model with TensorRT")
trt_graph = trt.create_inference_graph(
input_graph_def=frozen_graph,
outputs=output_names,
max_batch_size=1,
max_workspace_size_bytes=1 << 26,
precision_mode='FP16',
minimum_segment_size=2
)
print("Write optimized model to the file")
with open(DATA_DIR + MODEL + '_fp16_trt.pb', 'wb') as f:
f.write(trt_graph.SerializeToString())
在模型动物园的 ssd_mobilenet_v1_coco、ssd_mobilenet_v2_coco 和 ssd_inception_v2_coco 上进行了测试,它们的行为方式都相同 - 下载的 pb 文件在几秒钟内加载,TRT 优化 - 超过 10 分钟。怎么了?有没有人经历过同样的事情并且有任何提示如何解决它?