tensorflow自定义op_TensorRT加速tensorflow模型

import tensorflow as tf
import cv2
import base64
from tensorflow.python.compiler.tensorrt import trt_convert as trtsaved_model_dir &＃61; "./export_model_0126/1581080318"
output_saved_model_dir &＃61; "./convert_INT8_export_model"
fetch_names &＃61; ["strided_slice_256:0", "cond/Merge:0", "strided_slice_1:0", "ExpandDims:0","Const_39:0", "strided_slice_258:0", "Shape:0", "Cast:0","strided_slice_260:0", "combined_non_max_suppression/CombinedNonMaxSuppression:3","cond/Merge_1:0", "strided_slice_2:0"]class feed_dict_input_fn():def __init__(self, filename):self.filename &＃61; filenameself.content &＃61; []with open(self.filename) as f:for line in f:self.content.append(line.strip())self.index &＃61; 0def __call__(self, *args, **kwargs):data &＃61; open(self.content[self.index], &＃39;rb&＃39;).read()encode &＃61; base64.urlsafe_b64encode(data)encode &＃61; str(encode, encoding&＃61;&＃39;utf-8&＃39;)image &＃61; {"input:0": encode}# value &＃61; {"inputs": image}self.index &＃43;&＃61; 1return imageconverter &＃61; trt.TrtGraphConverter(input_saved_model_dir&＃61;saved_model_dir,precision_mode&＃61;trt.TrtPrecisionMode.INT8,use_calibration&＃61;True, is_dynamic_op&＃61;True, maximum_cached_engines&＃61;3)
feet_dict_input &＃61; feed_dict_input_fn("/home/admin-seu/TempData/sss/Master_work/data/test.list")
converter.convert()
converter.calibrate(fetch_names&＃61;fetch_names, num_runs&＃61;100, feed_dict_fn&＃61;feet_dict_input)
converter.save(output_saved_model_dir)

port tensorflow as tf
import numpy as np
import base64
import time# output_saved_model_dir &＃61; "./convert_export_model"
# output_saved_model_dir &＃61; "./export_model_0126/1581080318"
output_saved_model_dir &＃61; "./convert_INT8_export_model"
# output_saved_model_dir &＃61; "./convert_FP32_export_model"data &＃61; open("/home/admin-seu/TempData/test2017/000000258074.jpg", &＃39;rb&＃39;).read()
encode &＃61; base64.urlsafe_b64encode(data)
encode &＃61; str(encode, encoding&＃61;&＃39;utf-8&＃39;)with tf.Session() as sess:tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING],output_saved_model_dir)cur_graph &＃61; sess.graphnode_names &＃61; [tensor.name for tensor in sess.graph_def.node]output_tensors &＃61; []input_tensor &＃61; cur_graph.get_tensor_by_name("input:0")output_tensors.append(cur_graph.get_tensor_by_name("strided_slice_256:0"))output_tensors.append(cur_graph.get_tensor_by_name("strided_slice_260:0"))# for node_name in node_names:# if "input" in node_name:# print(node_name)# input_tensor &＃61; cur_graph.get_tensor_by_name(node_name)# if "scores_1" in node_name:# print(node_name)# output_tensors.append(cur_graph.get_tensor_by_name(node_name))# if "labels" in node_name:# print(node_name)# output_tensors.append(cur_graph.get_tensor_by_name(node_name))# if "all_ids" in node_name:# print(node_name)# output_tensors.append(cur_graph.get_tensor_by_name(node_name))# if "boxes_1" in node_name:# print(node_name)# output_tensors.append(cur_graph.get_tensor_by_name(node_name))output &＃61; sess.run(output_tensors, feed_dict&＃61;{input_tensor: encode})print(np.shape(output[0]))for i in range(10):output &＃61; sess.run(output_tensors, feed_dict&＃61;{input_tensor: encode})start &＃61; time.time()print(start)for i in range(100):output &＃61; sess.run(output_tensors, feed_dict&＃61;{input_tensor: encode})end &＃61; time.time()print(end)print(end - start)