88浏览
查看: 88|回复: 0

NVIDIA Jetson Nano 2GB 系列文章65:执行部署的 TensorRT 加速引擎

[复制链接]
qw1.jpg

很多开发人员在转换完 TensorRT 加速引擎之后,最后准备调用起来执行推理任务的时候,就遇到一些障碍。这个环节是需要开发人员自行撰写相关代码,去执行读入数据(前处理)、执行推理、显示结果(后处理)等工作,如下图最右边的部分。

qw2.jpg

这部分的麻烦之处,在于每个神经网络的结构不相同,并没有“通用”的代码可以适用于大部分的网络结构,需要针对指定神经网络去撰写对应的代码,最重要是需要清除这个模型的输入 (input bold) 与输出 (outpold) 的名称与张量结构。

本文以前面在 TAO 工具套件中使用的 ssd 神经网络为范例,提供基础的“前后处理”范例代码给读者参考,这是从 NVIDIA 中国区开发者社区所举办过多届 “Sky 黑客松”比赛中,所提供的开源内容中提取的重点,主要如下:

1、数据前处理:
  1.   def _preprocess_trt(img, shape=(300, 300)):
  2.     """TRT SSD推理前的数据前处理"""
  3.     img = cv2.resize(img, shape)
  4.     img = img.transpose((2, 0, 1)).astype(np.float32)
  5.     return img
复制代码



这里 “shape=(300,300)” 为张量的尺度,根据模型训练时的长宽两个变量,至于 transpose 里的 (2,0,1) 是固定的,不需调整。

2、数据后处理:


  1. def _postprocess_trt(img, output, conf_th, output_layout):
  2.     """TRT SSD推理后的结果的数据处理步骤."""
  3.     img_h, img_w, _ = img.shape
  4.     boxes, confs, clss = [], [], []
  5.     for prefix in range(0, len(output), output_layout):
  6.         index = int(output[prefix+0])
  7.         conf = float(output[prefix+2])
  8.         if conf < conf_th:
  9.             continue
  10.         x1 = int(output[prefix+3] * img_w)
  11.         y1 = int(output[prefix+4] * img_h)
  12.         x2 = int(output[prefix+5] * img_w)
  13.         y2 = int(output[prefix+6] * img_h)
  14.         cls = int(output[prefix+1])
  15.         boxes.append((x1, y1, x2, y2))
  16.         confs.append(conf)
  17.         clss.append(cls)
  18.     return boxes, confs, clss   # 返回标框坐标、置信度、类别
复制代码

这里最重要的 x1, y1,x2, y2 坐标值,必须根据 SSD 神经网络所定义的规范去进行修改,其他部分可以通用于大部分神经网络。

3、定义 TrtSSD 类封装运行 TRT SSD 所需的东西:

  1. class TrtSSD(object):
  2. # 加载自定义组建,如果TRT版本小于7.0需要额外生成flattenconcat自定义组件库
  3.     def _load_plugins(self):
  4.         if trt.__version__[0] < '7':
  5.             ctypes.CDLL("ssd/libflattenconcat.so")
  6.         trt.init_libnvinfer_plugins(self.trt_logger, '')
  7. #加载通过Transfer Learning Toolkit生成的推理引擎
  8.     def _load_engine(self):
  9.         TRTbin = 'ssd/TRT_%s.bin' % self.model  #请根据实际状况自行修改
  10.         with open(TRTbin, 'rb') as f, trt.Runtime(self.trt_logger) as runtime:
  11.             return runtime.deserialize_cuda_engine(f.read())
  12. #通过加载的引擎,生成可执行的上下文
  13.     def _create_context(self):
  14.         for binding in self.engine:
  15.             size = trt.volume(self.engine.get_binding_shape(binding)) * \
  16.                    self.engine.max_batch_size
  17. ##注意:这里的host_mem需要使用pagelocked memory,以免内存被释放
  18.             host_mem = cuda.pagelocked_empty(size, np.float32)
  19.             cuda_mem = cuda.mem_alloc(host_mem.nbytes)
  20.             self.bindings.append(int(cuda_mem))
  21.             if self.engine.binding_is_input(binding):
  22.                 self.host_inputs.append(host_mem)
  23.                 self.cuda_inputs.append(cuda_mem)
  24.             else:
  25.                 self.host_outputs.append(host_mem)
  26.                 self.cuda_outputs.append(cuda_mem)
  27.         return self.engine.create_execution_context()
  28. # 初始化引擎
  29.     def __init__(self, model, input_shape, output_layout=7):
  30.         self.model = model
  31.         self.input_shape = input_shape
  32.         self.output_layout = output_layout
  33.         self.trt_logger = trt.Logger(trt.Logger.INFO)
  34.         self._load_plugins()
  35.         self.engine = self._load_engine()
  36.         self.host_inputs = []
  37.         self.cuda_inputs = []
  38.         self.host_outputs = []
  39.         self.cuda_outputs = []
  40.         self.bindings = []
  41.         self.stream = cuda.Stream()
  42.         self.context = self._create_context()
  43. # 释放引擎,释放GPU显存,释放CUDA流
  44.     def __del__(self):
  45.         del self.stream
  46.         del self.cuda_outputs
  47.         del self.cuda_inputs
  48. # 利用生成的可执行上下文执行推理
  49.     def detect(self, img, conf_th=0.3):
  50.         img_resized = _preprocess_trt(img, self.input_shape)
  51.         np.copyto(self.host_inputs[0], img_resized.ravel())
  52.         # 将处理好的图片从CPU内存中复制到GPU显存
  53.         cuda.memcpy_htod_async(
  54.             self.cuda_inputs[0], self.host_inputs[0], self.stream)
  55.         # 开始执行推理任务
  56.         self.context.execute_async(
  57.             batch_size=1,
  58.             bindings=self.bindings,
  59.             stream_handle=self.stream.handle)
  60.         # 将推理结果输出从GPU显存复制到CPU内存
  61.         cuda.memcpy_dtoh_async(
  62.             self.host_outputs[1], self.cuda_outputs[1], self.stream)
  63.         cuda.memcpy_dtoh_async(
  64.             self.host_outputs[0], self.cuda_outputs[0], self.stream)
  65.         self.stream.synchronize()
  66.         output = self.host_outputs[0]
  67.         return _postprocess_trt(img, output, conf_th, self.output_layout)
复制代码


上面三个部分对不同神经网络都是不同的内容,如果要参考 YOLO 神经网络的对应内容,推荐参考 https://github.com/jkjung-avt/tensorrt_demos开源项目,里面有完整的 YOLOv3 与 YOLOv4 的详细内容。

本文的开源代码可以在此链接下载完整的内容与配套的工具。

https://pan.baidu.com/s/1fGLBnzqtnRNpfD3PbileOA  密码: 99et


您需要登录后才可以回帖 登录 | 立即注册

本版积分规则

为本项目制作心愿单
购买心愿单
心愿单 编辑
[[wsData.name]]

硬件清单

  • [[d.name]]
btnicon
我也要做!
点击进入购买页面
上海智位机器人股份有限公司 沪ICP备09038501号-4

© 2013-2022 Comsenz Inc. Powered by Discuz! X3.4 Licensed

mail