本文整理匯總了Python中tensorflow.contrib.tensorrt.calib_graph_to_infer_graph方法的典型用法代碼示例。如果您正苦於以下問題：Python tensorrt.calib_graph_to_infer_graph方法的具體用法？Python tensorrt.calib_graph_to_infer_graph怎麽用？Python tensorrt.calib_graph_to_infer_graph使用的例子？那麽, 這裏精選的方法代碼示例或許可以為您提供幫助。您也可以進一步了解該方法所在類tensorflow.contrib.tensorrt的用法示例。
在下文中一共展示了tensorrt.calib_graph_to_infer_graph方法的2個代碼示例，這些例子默認根據受歡迎程度排序。您可以為喜歡或者感覺有用的代碼點讚，您的評價將有助於係統推薦出更棒的Python代碼示例。
示例1: get_trt_graph_from_calib
# 需要導入模塊: from tensorflow.contrib import tensorrt [as 別名]
# 或者: from tensorflow.contrib.tensorrt import calib_graph_to_infer_graph [as 別名]
def get_trt_graph_from_calib(graph_name, calib_graph_def, output_dir):
"""Convert a TensorRT graph used for calibration to an inference graph."""
trt_graph = trt.calib_graph_to_infer_graph(calib_graph_def)
write_graph_to_file(graph_name, trt_graph, output_dir)
return trt_graph
################################################################################
# Run the graph in various precision modes.
################################################################################ 示例2: load_model
# 需要導入模塊: from tensorflow.contrib import tensorrt [as 別名]
# 或者: from tensorflow.contrib.tensorrt import calib_graph_to_infer_graph [as 別名]
def load_model(model, input_map=None):
# Check if the model is a model directory (containing a metagraph and a checkpoint file)
# or if it is a protobuf file with a frozen graph
model_exp = os.path.expanduser(model)
if (os.path.isfile(model_exp)):
print('Model filename: %s' % model_exp)
with gfile.FastGFile(model_exp,'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
#JJia TensorRT enable
print('TensorRT Enabled')
trt_graph = trt.create_inference_graph(input_graph_def=graph_def,
outputs=['embeddings:0'],
max_batch_size = 1,
max_workspace_size_bytes= 500000000, # 500MB mem assgined to TRT
precision_mode="FP16", # Precision "FP32","FP16" or "INT8"
minimum_segment_size=1
)
##trt_graph=trt.calib_graph_to_infer_graph(trt_graph)
#tf.import_graph_def(trt_graph, input_map=input_map, name='')
return trt_graph #"return graph_def" for trt disable, "return trt_graph" for trt enable
else:
print('Model directory: %s' % model_exp)
meta_file, ckpt_file = get_model_filenames(model_exp)
print('Metagraph file: %s' % meta_file)
print('Checkpoint file: %s' % ckpt_file)
saver = tf.train.import_meta_graph(os.path.join(model_exp, meta_file), input_map=input_map)
saver.restore(tf.get_default_session(), os.path.join(model_exp, ckpt_file))
#JJia TensorRT enable
print('TensorRT Enabled', 1<<20)
frozen_graph = tf.graph_util.convert_variables_to_constants(
tf.get_default_session(),
tf.get_default_graph().as_graph_def(),
output_node_names=["embeddings"])
for node in frozen_graph.node:
if node.op == 'RefSwitch':
node.op = 'Switch'
elif node.op == 'AssignSub':
node.op = 'Sub'
if 'use_locking' in node.attr: del node.attr['use_locking']
trt_graph = trt.create_inference_graph(
input_graph_def=frozen_graph,
outputs=["embeddings"],
max_batch_size = 1,
max_workspace_size_bytes= 1 << 20,
precision_mode="FP16",
minimum_segment_size=1)
#tf.import_graph_def(trt_graph,return_elements=["embeddings:0"])
return trt_graph #"return frozen_graph" for trt disable, "return trt_graph" for trt enable