如何从保存的模型中导出RT图

时间:2019-08-25 07:59:20

标签: python tensorflow tensorrt

我想将图形导出到RT进行推理。如何使用保存的模型来实现。用于保存的模型的推断代码如下:

import tensorflow as tf
import cv2

img = cv2.imread("test.jpg")
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_rgb = np.expand_dims(img, 0)

predict_fn = tf.contrib.predictor.from_saved_model("./saved_model")
output_data = predict_fn({"inputs": img_rgb})
print(output_data)  # detector output dictionary

我不想使用frozen_inference_garph.pb。我想使用saved_model/saved_model.pb进行推理(因为它的准确性更高)。

1 个答案:

答案 0 :(得分:0)

命令行工具saved_model_cli对我有用。

saved_model_cli convert \
--dir "/home/yilrr/tf-serving/faster-rcnn/saved_model/versions/1" \
--output_dir "/home/yilrr/tf-serving/trt-frcnn" \
--tag_set serve \
tensorrt --precision_mode FP32 --max_batch_size 32 --is_dynamic_op True

以下是推理代码:

import tensorflow as tf
import cv2
import numpy as np

IMAGE_PATH = "/Users/vedanshu/test.jpg"
image = cv2.imread(IMAGE_PATH)

input_names = ['image_tensor']

def draw_label(image, point, label, font=cv2.FONT_HERSHEY_SIMPLEX,
               font_scale=0.5, thickness=2):
    size = cv2.getTextSize(label, font, font_scale, thickness)[0]
    x, y = point
    cv2.rectangle(image, (x, y - size[1]),
                  (x + size[0], y), (255, 0, 0), cv2.FILLED)
    cv2.putText(image, label, point, font, font_scale,
                (255, 255, 255), thickness)

detection_graph1 = tf.Graph()

with detection_graph1.as_default():
    tf_sess1 = tf.Session(graph=detection_graph1)
    model = tf.saved_model.loader.load(tf_sess1, ["serve"], "/Users/vedanshu/ckpt/saved_model")

tf_input = tf_sess1.graph.get_tensor_by_name(input_names[0] + ':0')
tf_scores = tf_sess1.graph.get_tensor_by_name('detection_scores:0')
tf_boxes = tf_sess1.graph.get_tensor_by_name('detection_boxes:0')
tf_classes = tf_sess1.graph.get_tensor_by_name('detection_classes:0')
tf_num_detections = tf_sess1.graph.get_tensor_by_name('num_detections:0')
scores, boxes, classes, num_detections = tf_sess1.run([tf_scores, tf_boxes, tf_classes, tf_num_detections], feed_dict={tf_input: image[None, ...]})

boxes = boxes[0]  # index by 0 to remove batch dimension
scores = scores[0]
classes = classes[0]
num_detections = int(num_detections[0])

boxes_pixels = []
for i in range(num_detections):
    # scale box to image coordinates
    box = boxes[i] * np.array([image.shape[0],
                               image.shape[1], image.shape[0], image.shape[1]])
    box = np.round(box).astype(int)
    boxes_pixels.append(box)

boxes_pixels = np.array(boxes_pixels)

for i in range(num_detections):
    box = boxes_pixels[i]
    box = np.round(box).astype(int)
    image = cv2.rectangle(image, (box[1], box[0]), (box[3], box[2]), (0, 255, 0), 2)
    label = "{}:{:.2f}".format(int(classes[i]), scores[i])
    draw_label(image, (box[1], box[0]), label)

cv2.imwrite("/Users/vedanshu/out.jpg", image)