pip install nvidia-pyindex
pip install nvidia-tensorrt
封装过程
- 构建阶段
import tensorrt as trt
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
- 创建网络
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
- 反序列化加载engine
runtime = trt.Runtime(logger)
engine = runtime.deserialize_cuda_engine(serialized_engine)
- 执行推理
context = engine.create_execution_context()
context.execute_async_v2(buffers, stream_ptr)
- 整体封装
import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
trt.init_libnvinfer_plugins(TRT_LOGGER, '')
def GiB(val):
return val * 1 << 30
# Simple helper data class that's a little nicer to use than a 2-tuple.
class HostDeviceMem(object):
def __init__(self, host_mem, device_mem):
self.host = host_mem
self.device = device_mem
def __str__(self):
return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)
def __repr__(self):
return self.__str__()
# Allocates all buffers required for an engine, i.e. host/device inputs/outputs.
def allocate_buffers(engine):
inputs = []
outputs = []
bindings = []
stream = cuda.Stream()
for binding in engine:
size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
dtype = trt.nptype(engine.get_binding_dtype(binding))
# Allocate host and device buffers
host_mem = cuda.pagelocked_empty(size, dtype)
device_mem = cuda.mem_alloc(host_mem.nbytes)
# Append the device buffer to device bindings.
bindings.append(int(device_mem))
# Append to the appropriate list.
if engine.binding_is_input(binding):
inputs.append(HostDeviceMem(host_mem, device_mem))
else:
outputs.append(HostDeviceMem(host_mem, device_mem))
return inputs, outputs, bindings, stream
# This function is generalized for multiple inputs/outputs.
# inputs and outputs are expected to be lists of HostDeviceMem objects.
def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):
# Transfer input data to the GPU.
[cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
# Run inference.
context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)
# Transfer predictions back from the GPU.
[cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
# Synchronize the stream
stream.synchronize()
# Return only the host outputs.
return [out.host for out in outputs]
# This function is generalized for multiple inputs/outputs for full dimension networks.
# inputs and outputs are expected to be lists of HostDeviceMem objects.
def do_inference_v2(context, bindings, inputs, outputs, stream):
# Transfer input data to the GPU.
[cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
# Run inference.
context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
# Transfer predictions back from the GPU.
[cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
# Synchronize the stream
stream.synchronize()
# Return only the host outputs.
return [out.host for out in outputs]
# The Onnx path is used for Onnx models.
def build_engine_onnx(model_file):
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(EXPLICIT_BATCH)
config = builder.create_builder_config()
parser = trt.OnnxParser(network, TRT_LOGGER)
config.max_workspace_size = GiB(1)
# Load the Onnx model and parse it in order to populate the TensorRT network.
with open(model_file, 'rb') as model:
if not parser.parse(model.read()):
print('ERROR: Failed to parse the ONNX file.')
for error in range(parser.num_errors):
print(parser.get_error(error))
return None
return builder.build_engine(network, config)
# The Caffe path is used for Caffe2 models.
def build_engine_caffe(model_file, deploy_file, dtype, mark_output: list):
# You can set the logger severity higher to suppress messages (or lower to display more messages).
with trt.Builder(
TRT_LOGGER) as builder, builder.create_network() as network, builder.create_builder_config() as config, trt.CaffeParser() as parser:
# Workspace size is the maximum amount of memory available to the builder while building an engine.
# It should generally be set as high as possible.
config.max_workspace_size = GiB(1)
# Load the Caffe model and parse it in order to populate the TensorRT network.
# This function returns an object that we can query to find tensors by name.
model_tensors = parser.parse(deploy=deploy_file, model=model_file, network=network, dtype=dtype)
# For Caffe, we need to manually mark the output of the network.
# Since we know the name of the output tensor, we can find it in model_tensors.
for layer in mark_output:
network.mark_output(model_tensors.find(layer))
return builder.build_engine(network, config)
class Model:
def __init__(self, engine_file=None):
"""
:param engine_file: engine 文件
"""
self.engine = None
if engine_file:
self.load_engine_file(engine_file=engine_file)
self.inputs, self.outputs, self.bindings, self.stream, self.context = None, None, None, None, None
def load_engine_file(self, engine_file):
"""
加载trt engine文件
:param engine_file:
:return:
"""
runtime = trt.Runtime(TRT_LOGGER)
with open(engine_file, "rb") as f:
engine = runtime.deserialize_cuda_engine(f.read())
self.engine = engine
def build_engine_by_onnx(self, model_file="ResNet50.onnx"):
"""
通过onnx构建engine文件
:param model_file:
:return:
"""
self.engine = build_engine_onnx(model_file)
def build_engine_by_caffe(self, model_file="ResNet50_fp32.caffemodel", deploy_file="ResNet50_N2.prototxt",
dtype=trt.float32, mark_output: list = []):
"""
通过caffe文件构建engine
:param model_file:
:param deploy_file:
:param dtype:
:param mark_output:
:return:
"""
self.engine = build_engine_caffe(model_file, deploy_file, dtype, mark_output=mark_output)
def allocate_buffers(self):
"""
分配缓存
:return:
"""
self.inputs, self.outputs, self.bindings, self.stream = allocate_buffers(self.engine)
def create_execution_context(self):
"""
生成执行器
:return:
"""
self.context = self.engine.create_execution_context()
def do_inference(self, inputs):
"""
推理
:param inputs:
:return:
"""
self.inputs[0].host[:] = inputs[:]
return do_inference(context=self.context, bindings=self.bindings, inputs=self.inputs, outputs=self.outputs,
stream=self.stream)
def do_inference_v2(self, inputs):
"""
推理
:param inputs:
:return:
"""
self.inputs[0].host[:] = inputs[:]
return do_inference_v2(context=self.context, bindings=self.bindings, inputs=self.inputs, outputs=self.outputs,
stream=self.stream)
- 测试
自行准备模型engine文件
import numpy as np
import cv2
from infer import Model
# 查找网络的所有输出层
def getOutputLayers(prototxt_path):
# 打开prototxt
f = open(prototxt_path)
# 按行读取
line = f.readline()
# 记录bottom的名字
bottom = []
# 记录top的名字
top = []
# 记录每个bottom和top的行号,解决输出层因为inplace无法正确获取的问题
cnt = 0
while line:
line = line.replace(" ", "")
line = line.replace('"', "")
pos = 0
flag = ""
for j in range(len(line)):
if line[j] == ':':
pos = j + 1
break
else:
flag += line[j]
name = ""
if flag != "":
for j in range(pos, len(line) - 1):
name += line[j]
if flag == "bottom":
bottom.append([name, cnt])
if flag == "top":
top.append([name, cnt])
line = f.readline()
cnt = cnt + 1
f.close()
ans = []
for i in range(len(top)):
now = top[i][0]
now_cnt = top[i][1]
flag = True
# 只有当前top的名字在bottom中找不到,或者只能找到一个距离最近的(inplace情况),才是输出层
for j in range(len(bottom)):
if bottom[j][0] == now:
if abs(bottom[j][1] - now_cnt) > 1:
flag = False
if (flag):
ans.append(now)
# 返回输出层名字列表,对于ssd就是['detection_out', 'keep_count']
return ans
mark_output = []
layers = getOutputLayers("head_detection.prototxt")
for layer in layers:
mark_output.append(layer)
model = Model(engine_file="head_detection.engine")
model.build_engine_by_caffe(model_file="./head_detection.caffemodel", deploy_file="head_detection.prototxt",
mark_output=mark_output)
model.allocate_buffers()
model.create_execution_context()
print(model.engine)
print(model.bindings)
img = cv2.imread("./1.jpg")
cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
cv2.resize(img, (512, 512))
data = np.array(img).flatten()
res = model.do_inference(data)[0].reshape((1, 1, 50, 7))
print(res)
print(res.shape)
- 效果
root@d697bba88e88:/workspace/build/models# python3 test_trt.py
[TensorRT] WARNING: Current optimization profile is: 0. Please ensure there are no enqueued operations pending in this context prior to switching profiles
<tensorrt.tensorrt.ICudaEngine object at 0x7f9cddc21930>
[140308738211840, 140309111512576, 140309111514112]
(598, 900, 3)
resieze = (512, 512, 3)
[ 98 87 81 ... 221 221 221]
(786432,)
[[[[0.00000000e+00 1.00000000e+00 7.01323450e-01 2.53085375e-01
0.00000000e+00 2.63810396e-01 1.81251913e-02]
[0.00000000e+00 1.00000000e+00 6.86035991e-01 2.68468022e-01
0.00000000e+00 2.79185116e-01 1.83350369e-02]
[0.00000000e+00 1.00000000e+00 6.22933388e-01 7.85305619e-01
0.00000000e+00 7.94009924e-01 1.48523413e-02]
[0.00000000e+00 1.00000000e+00 6.11043394e-01 9.83310580e-01
0.00000000e+00 9.96181369e-01 1.94536522e-02]
[0.00000000e+00 1.00000000e+00 6.08991027e-01 7.70242155e-01
0.00000000e+00 7.79535949e-01 1.44858146e-02]
[0.00000000e+00 1.00000000e+00 6.05013907e-01 2.83096880e-01
欢迎分享,转载请注明来源:内存溢出
评论列表(0条)