Python api构建tensorrt加速模型的步骤详解

import tensorrt as trt

EXPLICIT_BATCH = 1 << (int)
(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
with builder = trt.Builder(TRT_LOGGER) as builder,
	builder.create_network(EXPLICIT_BATCH) as network, trt.OnnxParser(network,
 TRT_LOGGER) as parser:
with open(model_path, 'rb') as model:
parser.parse(model.read())

builder.max_batch_size = max_batch_size
builder.max_workspace_size = 1 << 20

#序列化模型到模型流
serialized_engine = engine.serialize()
#反序列化模型流去执行推理，反序列化需要创建一个运行时对象
with trt.Runtime(TRT_LOGGER) as runtime:
	engine = runtime.deserialize_cuda_engine(serialized_engine)
#也可以将序列化模型write
with open(“sample.engine”, “wb”) as f:
	f.write(engine.serialize())
#然后再读出来进行反序列化
with open(“sample.engine”, “rb”) as f, trt.Runtime(TRT_LOGGER) as runtime:
	engine = runtime.deserialize_cuda_engine(f.read())

# 为输入和输出分配一些主机和设备缓冲区:
#确定尺寸并创建页面锁定内存缓冲区
h_input = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(0)),dtype=np.float32)
h_output =cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(1)),dtype=np.float32)
#为输入和输出分配设备内存
d_input = cuda.mem_alloc(h_input.nbytes)
d_output = cuda.mem_alloc(h_output.nbytes)
#创建一个流，在其中复制输入/输出并运行推断
stream = cuda.Stream()

# 创建一些空间来存储中间激活值，因为engine保存了network定义和训练时的参数，这些都是构建的上下文执行的。
with engine.create_execution_context() as context:
 	# 输入数据传入GPU
	cuda.memcpy_htod_async(d_input, h_input, stream)
 	# 执行推理.
 	context.execute_async(bindings=[int(d_input), int(d_output)],
 	stream_handle=stream.handle)
 	# 将推理后的预测结果从GPU上返回.
 	cuda.memcpy_dtoh_async(h_output, d_output, stream)
 	# 同步流
 	stream.synchronize()
 	# 返回主机输出 
	return h_output