在上一篇文章中,对构建TensorRT引擎进行了介绍,本文将详细介绍如何使用Python API通过TensorRT引擎加速,实现单幅图像的推断
Workflow
单幅图像的推理过程大致如下
-
读入图像文件,并进行预处理
-
加载TensorRT引擎文件并进行反序列化,创建执行上下文context
-
根据引擎申请输入输出Buffers,并进行绑定
-
调用推断API进行推断
-
输出解码
-
画框并显示
实现
首先是导入需要的包
import cv2
import time
import numpy as np
import logging
import argparse
import tensorrt as trt
cv2用于进行图像文件加载、图像预处理以及画框,time用于进行耗时统计,numpy用于进行前后处理中的张量操作,argparse用于构造一个参数解析器,tensorrt是TensorRT封装API包
预先设置锚框和锚框掩码,这里需要和训练的模型输出对应上,训练中使用的9个锚框分为了3个尺度,掩码用于标记锚框属于哪个尺度。另外训练过程的锚框数值是用的归一化值还是绝对值必须清楚,我这里锚框的归一化是在解码函数中处理的,因此这里放的是绝对值
anchors = np.array([(10, 13), (16, 30), (33, 23), (30, 61), (62, 45),
(59, 119), (11