У меня очень долгое время загрузки графиков TensorFlow, оптимизированных с помощью TensorRT. Неоптимизированные загружаются быстро, но загрузка оптимизированных занимает более 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 минут. Что случилось? Кто-нибудь испытал то же самое и есть подсказки, как это исправить?