如果还未配置TensorRT,请看这篇博文:Win11下TensorRT环境部署
这里使用TensorRT对Yolov5进行部署流程比较固定:先将pt模型转换为onnx,再将onnx模型转为engine,所以在执行export.py时要将onnx、engine给到include。
PT模型地址:
python path/to/export.py --weights yolov5s.pt --include torchscript onnx coreml saved_model pb tflite tfjs
我们可以看到:
这里介绍一下其中每个参数的含义:
为什么要把模型和输入数据转移到CPU上?
CPU 是所有计算环境中最通用和兼容的硬件。当你将一个模型导出为一个通用格式,如 ONNX,目的通常是确保它可以在不同的环境和硬件上运行。使用 CPU 可以确保最大程度的兼容性,因为不是所有的环境都有 GPU 支持。这有助于确保模型能够在各种不同的硬件和软件环境中一致和可靠地工作。
如果要精简onnx,就可以将simplify设置为True。设置为True,就会调用onnxsim来对原来onnx去除不必要的op操作,也叫去除胶水,在使用之前需要安装onnx-simplifer。
精简完之后可以把简化前和简化后的模型放进netron看优化了哪些地方。netron 就是一个网络结构的可视化神器。我们可以通过它来查看网络的结构。因为这里的模型太大,会占用较长的篇幅,所以本文就不展示了。
构建完onnx后,就开始从onnx转为engine。在TensorRT上主要存在以下几个对象:
LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...')
assert onnx.exists(), f'failed to export ONNX file: {onnx}'
f = file.with_suffix('.engine') # TensorRT engine file
logger = trt.Logger(trt.Logger.INFO)
if verbose:
logger.min_severity = trt.Logger.Severity.VERBOSE
builder = trt.Builder(logger)
config = builder.create_builder_config()
config.max_workspace_size = workspace * 1 << 30
# config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30) # fix TRT 8.4 deprecation notice
flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
network = builder.create_network(flag)
parser = trt.OnnxParser(network, logger)
if not parser.parse_from_file(str(onnx)):
raise RuntimeError(f'failed to load ONNX file: {onnx}')
inputs = [network.get_input(i) for i in range(network.num_inputs)]
outputs = [network.get_output(i) for i in range(network.num_outputs)]
for inp in inputs:
LOGGER.info(f'{prefix} input "{inp.name}" with shape{inp.shape} {inp.dtype}')
for out in outputs:
LOGGER.info(f'{prefix} output "{out.name}" with shape{out.shape} {out.dtype}')
if dynamic:
if im.shape[0] <= 1:
LOGGER.warning(f'{prefix} WARNING ?? --dynamic model requires maximum --batch-size argument')
profile = builder.create_optimization_profile()
for inp in inputs:
profile.set_shape(inp.name, (1, *im.shape[1:]), (max(1, im.shape[0] // 2), *im.shape[1:]), im.shape)
config.add_optimization_profile(profile)
LOGGER.info(f'{prefix} building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine as {f}')
if builder.platform_has_fast_fp16 and half:
config.set_flag(trt.BuilderFlag.FP16)
with builder.build_engine(network, config) as engine, open(f, 'wb') as t:
t.write(engine.serialize())
return f, None
使用engine推理和使用pt推理的流程是大同小异的,同样是使用detect.py。区别就在于model = DetectMultiBackend(weights, device=device, dnn=dnn)中的weight是pt模型还是onnx模型还是engine模型
进入到DetectMultiBackend这个类中查看,直接跳转【ctrl + click】到构建engine的部分:
elif engine: # TensorRT
LOGGER.info(f'Loading {w} for TensorRT inference...')
import tensorrt as trt # https://developer.nvidia.com/nvidia-tensorrt-download
check_version(trt.__version__, '7.0.0', hard=True) # require tensorrt>=7.0.0
if device.type == 'cpu':
device = torch.device('cuda:0')
Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr'))
logger = trt.Logger(trt.Logger.INFO)
with open(w, 'rb') as f, trt.Runtime(logger) as runtime:
model = runtime.deserialize_cuda_engine(f.read())
context = model.create_execution_context()
bindings = OrderedDict()
output_names = []
fp16 = False # default updated below
dynamic = False
for i in range(model.num_bindings):
name = model.get_binding_name(i)
dtype = trt.nptype(model.get_binding_dtype(i))
if model.binding_is_input(i):
if -1 in tuple(model.get_binding_shape(i)): # dynamic
dynamic = True
context.set_binding_shape(i, tuple(model.get_profile_shape(0, i)[2]))
if dtype == np.float16:
fp16 = True
else: # output
output_names.append(name)
shape = tuple(context.get_binding_shape(i))
im = torch.from_numpy(np.empty(shape, dtype=dtype)).to(device)
bindings[name] = Binding(name, dtype, shape, im, int(im.data_ptr()))
binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
batch_size = bindings['images'].shape[0] # if dynamic, this is instead max batch size
在进行前向推理时,就会调用DetectMultiBackend的forward方法:
其中y返回的就是推理结果。这个结果会返回到detect.py中,当detect.py捕获到结果之后,就正常走目标检测的后处理流程了。
–weight填写 yolov5s.onnx
–include 填写 engine
注:如果输入.onnx生成engine报错的话可以直接在–weight填我们需要优化的权重
–include 直接填写engine,因为直接生成engine的过程中也会升成一个.onnx。
参考链接:
https://zhuanlan.zhihu.com/p/607601799
https://blog.csdn.net/m0_64524798/article/details/129187608