1

我想使用 tensorrt 来加速 tensorflow 模型推断。但是这样做时我没有得到任何结果，以下是我的代码。返回框和分数均为无。我是否以错误的方式使用 TensorRT 和 tensorflow？

import tensorflow as tf
import tensorflow.contrib.tensorrt as trt
import os
from PIL import Image
import numpy as np

with tf.Session() as sess:
    with tf.gfile.GFile('model/ssd_inceptionv2.pb',  'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    trt_graph = trt.create_inference_graph(
            input_graph_def=graph_def,
            outputs=["detection_boxes", "detection_classes"])
    output_node = tf.import_graph_def(
            trt_graph,
            return_elements=["detection_boxes", "detecction_classes"], name='')
    input = tf.get_default_graph().get_tensor_by_name('image_tensor:0')

    for root, dirs, files in os.walk('dataset/test'):
        for f in files:
            if not f.endswith('.jpg'):
                continue
            print f
            image = Image.open(os.path.join(root, f))
            print image
            image = image.resize((300, 300), Image.ANTIALIAS)
            im_width, im_height = image.size
            image = np.array(image.getdata()).reshape((im_height, im_width, 3)).astype(np.uint8)
            image = np.expand_dims(image, axis=0)
            boxes, scores = sess.run(output_node, feed_dict={input: image})
            boxes, scores = np.squeeze(boxes), np.squeeze(scores)
            print boxes, scores
            for i, s in enumerate(scores):
                if s > 0.5:
                    print boxes[i], s
4

1 回答 1

0

问题出在 output_node 的定义中。return_elements 应定义为return_elments=["detection_boxes:0", "detection_classes:0"]

于 2019-07-18T11:18:11.710 回答