tensorflow - Yolov2 Frozen Graph To Tensorrt 图

Question

我使用以下代码将 yolov2 冻结图转换为 tftrt 图。

OUTPUT_NAME = ["models/convolutional23/BiasAdd"]

# read Tensorflow frozen graph
with gfile.FastGFile('./yolov2_frozen-graph.pb', 'rb') as tf_model:
   tf_graphf = tensorflow.GraphDef()
   tf_graphf.ParseFromString(tf_model.read())

# convert (optimize) frozen model to TensorRT model
trt_graph = trt.create_inference_graph(input_graph_def=tf_graphf, outputs=OUTPUT_NAME, max_batch_size=1, max_workspace_size_bytes=2 * (10 ** 9), precision_mode="FP32")

# write the TensorRT model to be used later for inference
with gfile.FastGFile("Yolo_TensorRT_modelFP16.pb", 'wb') as f:
   f.write(trt_graph.SerializeToString())
print("TensorRT model is successfully stored!")

然后，我使用以下代码运行推理。

with tf.Session() as sess:
    img = cv2.imread("image3.jpg")
    img = cv2.resize(img, (608, 608))
    
    # read TensorRT frozen graph
    with gfile.FastGFile('Yolo_TensorRT_modelFP16.pb', 'rb') as trt_model:
        trt_graph = tf.GraphDef()
        trt_graph.ParseFromString(trt_model.read())

    # obtain the corresponding input-output tensor
    tf.import_graph_def(trt_graph, name='')
    input = sess.graph.get_tensor_by_name('models/net1:0')
    output = sess.graph.get_tensor_by_name('models/convolutional23/BiasAdd:0')

    for i in range(100):
        start = time.time()
        # perform inference
        sess.run(output, feed_dict={input: [np.asarray(img)]})
        end =  time.time() - start
        print("infernce time: ", end)

因此，即使在我对 FP16 yolov2 冻结 tftrt 图进行推理之后，它的性能也与普通 yolov2 冻结图完全相同。你能告诉我我必须做些什么来提高 tftrt 图的性能吗？

标签： tensorflownvidiayolotensorrtnvidia-jetson