Skip to content
/ tensorRT Public

TensorRT-7 Network Lib 包括常用目标检测、关键点检测、人脸检测、OCR等 可训练自己数据

525 stars 111 forks Branches Tags Activity
Star
Notifications You must be signed in to change notification settings

Syencil/tensorRT

BranchesTags

Repository files navigation

TensorRT-7 Network Lib

Introduction

Python ===> Onnx ===> tensorRT ===> .h/.so
支持FP32，FP16，INT8量化。支持serialize，deserialize
基于线程池实现多线程并发，提升预处理和后处理的速度
重写或融合部分Opencv算子，提升Cache使用率以及避免不必要的扫描操作
支持infer时GPU和CPU端异步进行实现延迟隐藏
支持剪枝、蒸馏、量化、换轻量级backbone
推荐搭配https://github.com/Syencil/mobile-yolov5-pruning-distillation使用

Model Zoo

Model Training git Infer Time Total Time
Yolov5x https://github.com/ultralytics/yolov5
https://github.com/Syencil/mobile-yolov5-pruning-distillation		 32.5ms 58ms
PANNet(Pse++) https://github.com/WenmuZhou/PAN.pytorch 18.5ms 45ms
PSENet https://github.com/WenmuZhou/PSENet.pytorch 22ms 48ms
Yolov3 https://github.com/YunYang1994/tensorflow-yolov3 14.5ms 29.5ms
Retinaface https://github.com/biubug6/Pytorch_Retinaface
https://github.com/Syencil/Pytorch_Retinaface		 2.3ms 12.3ms
Retinanet mmdetection + configs/nas_fpn/retinanet_r50_fpn_crop640_50e_coco.py 22.9ms 333ms
Fcos mmdetection + configs/fcos/fcos_r50_caffe_fpn_4x4_1x_coco.py - -
ResNet - - -
Hourglass https://github.com/Syencil/Keypoints 28ms 37ms
SimplePose https://github.com/microsoft/human-pose-estimation.pytorch 3ms 7ms
  • 测试环境为Tesla P40 + 4个CPU线程。

Quick Start

Code -> Onnx

git Convert
tensorflow https://github.com/onnx/tensorflow-onnx python -m tf2onnx.convert
pytorch - torch.onnx.export(model, img, weights, verbose=False, opset_version=11, input_names=['images'], output_names=['output'])
Onnx onnx-simplifier python3 -m onnxsim in.onnx out.onnx

C++

mkdir build
cd build
cmake -DCMAKE_BUILD_TYPE=Release ..
make + project_lib
make + project_name
./bin/project_name

Tips

  • Onnx必须指定为输入全尺寸，再实际中trt也不存在理想上的动态输入，所以必须在freeze阶段指明输入大小。
  • 构建新项目时，通常只需要需要继承TensorRT类下面的DetectionTRT/SegmentationTRT/KeypointTRT类。只需要实现postProcess就可以用了。上层暴露出来的接口为initSession和predOneImage两个方法，方便调用。
  • 由于ONNX和TRT目前算子实现的比较完善，大多数时候只需要实现相应后处理即可，针对特定算子通常可以再python代码中用一些trick进行替换，实在不行可以考虑自定义plugin
  • 关于CHW和HWC的数据格式
    • CHW: 对于GPU更优。使用CUDA做infer或者后处理的话，由于硬件DRAM的原因，CHW可以保证线程是以coalescing的方式读取。具体性能对比参考Programming_Massively_Parallel_Processors
    • HWC: 对于CPU更优。使用CPU进行处理的时候，HWC格式可以保证单个线程处理的数据具有连续的内存地址。而CPU缓存具有空间局部性，这样能极大的提升效率。
    • 综上：如果后处理使用CPU进行decode，建议在onnx输出HWC格式，如果使用GPU进行decode，建议在onnx输出CHW格式。对于输入端则没有具体测试，主要是相信tensorflow虽然按照之前的HWC格式，但是在操作中肯定也是做了优化

Darknet

简介

  • 位置：yolov3_darknet_main.cpp

注意事项

利用pytorch-yolov4将darknet模型转换成onnx之后使用

SimplePose

简介

注意事项

  • 转出onnx之后，在解析onnx时，需要将tmp Cuda的空间设置大一点，不然解析deconv的时候会报错。

StreamProcess

简介

  • 位置：stream_main.cpp
  • 此项目为基于yolov5的GPU和CPU端分离之后进行延迟隐藏的简单demo
  • 以对视频进行推理和渲染为基础示例，可以自由更改或重写preFunc和postFunc来实现不同的需求

PanNet (PseNet V2)

简介

注意事项

  • pan和pse代码其实高度相似，导出的方法可以参考PseNet也可以参考我fork后改的代码。
  • pan网络中转出onnx的结果是没有经过sigmoid的(尝试一下加在后处理)
  • sigmoid在CPU中计算耗时比较大，可以参考fast-sigmoid-algorithm。 CPU上性能对比结果100000 times sigmoid ==> 2.81878ms fast sigmoid ==> 0.589737ms，而GPU上两者差异忽略不记。
    fast_sigmoid(x) = (x / (1 + |x|)) * 0.5 + 0.5

PseNet

简介

注意事项

  • torch转onnx的代码可以加在predict.py中，只需要在Pytorch_model这个类里面加一个成员函数即可
    def export(self, onnx_path, input_size):
        assert isinstance(input_size, list) or isinstance(input_size, tuple)
        self.net.export = True
        img = torch.zeros((1, 3, input_size[0], input_size[1])).to(self.device)
        with torch.no_grad():
            torch.onnx.export(self.net, img, onnx_path, verbose=True, opset_version=11, export_params=True, do_constant_folding=True)
        print("Onnx Simplify...")
        os.system("python3 -m onnxsim {} {}".format(onnx_path, onnx_path))
        print('Export complete. ONNX model saved to %s\nView with https://github.com/lutzroeder/netron' % onnx_path)
  • 为了方便trt的处理，我把sigmoid加入到了torch的代码中。在models/model.py中修改PSENet的forward代码，同时__init__中加入成员变量export=False来控制
        if self.export:
            x = torch.sigmoid(x)
        return x

Yolov5

简介

注意事项

  • trt的decode针对的是BxHxWxAC的格式（方便按height方向并行化以及其他嵌入式接入）。原版yolov5导出的onnx是BxAxHxWxC，需要在models/yolo.py第28行改为
            if self.export:
                x[i] = x[i].permute(0, 2, 3, 1).contiguous()
            else:
                x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()

RetinaFace

简介

注意事项

  • 执行convert_to_onnx.py的时候需要更改opset_version=11，verbose=True
  • 因为项目不需要关键点，所以把landmark的decode部分去掉了
  • 直接使用阈值0.6（原版0.02 + topK）过滤然后接NMS
  • 支持多线程操作

Yolov3

简介

注意事项

  • 训练部分同原版git相同，主要在freeze的时候使用了固定尺寸输入，并修改了python中decode的实现方法。修改为core/yolov3.py增加一个decode_full_shape类方法
    def decode_full_shape(self, conv_output, anchors, stride):
        """
        return tensor of shape [batch_size, output_size, output_size, anchor_per_scale, 5 + num_classes]
               contains (x, y, w, h, score, probability)
        """
        conv_shape = conv_output.get_shape().as_list()
        batch_size = conv_shape[0]
        output_size = conv_shape[1]
        anchor_per_scale = len(anchors)

        conv_output = tf.reshape(conv_output, (batch_size, output_size, output_size, anchor_per_scale, 5 + self.num_class), name="reshape")

        conv_raw_dxdy = conv_output[:, :, :, :, 0:2]
        conv_raw_dwdh = conv_output[:, :, :, :, 2:4]
        conv_raw_conf = conv_output[:, :, :, :, 4:5]
        conv_raw_prob = conv_output[:, :, :, :, 5:]

        y_np = np.tile(np.arange(output_size, dtype=np.int32)[..., np.newaxis], [1, output_size])
        x_np = np.tile(np.arange(output_size, dtype=np.int32)[np.newaxis, ...], [output_size, 1])

        xy_grid_np = np.concatenate([np.reshape(x_np, [np.shape(x_np)[0], np.shape(x_np)[1], 1]), np.reshape(y_np, [np.shape(y_np)[0], np.shape(y_np)[1], 1])], axis=2)
        xy_grid_np = np.tile(np.reshape(xy_grid_np, [1, np.shape(xy_grid_np)[0], np.shape(xy_grid_np)[1], 1, np.shape(xy_grid_np)[2]]), [batch_size, 1, 1, anchor_per_scale, 1])

        anchor_np = np.tile(np.reshape(anchors, [1, 1, 1, -1]), [batch_size, output_size, output_size, 1])

        xy_grid = tf.constant(xy_grid_np, dtype=tf.float32)
        stride_tf = tf.constant(shape=[batch_size, output_size, output_size, anchor_per_scale * 2], value=stride, dtype=tf.float32)
        anchor_tf = tf.constant(anchor_np, dtype=tf.float32)

        pred_xy = tf.sigmoid(conv_raw_dxdy)
        pred_wh = tf.exp(conv_raw_dwdh)

        pred_xy = tf.reshape(pred_xy, [batch_size, output_size, output_size, anchor_per_scale * 2])
        pred_wh = tf.reshape(pred_wh, [batch_size, output_size, output_size, anchor_per_scale * 2])
        xy_grid = tf.reshape(xy_grid, [batch_size, output_size, output_size, anchor_per_scale * 2])

        pred_xy = tf.add(pred_xy, xy_grid)
        pred_xy = tf.multiply(pred_xy, stride_tf)
        pred_wh = tf.multiply(pred_wh, anchor_tf)
        pred_wh = tf.multiply(pred_wh, stride_tf)

        pred_xy = tf.reshape(pred_xy, [batch_size, output_size, output_size, anchor_per_scale, 2])
        pred_wh = tf.reshape(pred_wh, [batch_size, output_size, output_size, anchor_per_scale, 2])

        pred_xywh = tf.concat([pred_xy, pred_wh], axis=4)

        pred_conf = tf.sigmoid(conv_raw_conf)
        pred_prob = tf.sigmoid(conv_raw_prob)

        return tf.concat([pred_xywh, pred_conf, pred_prob], axis=4, name="decode")
  • NMS代码采用faster-rcnn中的NMS。改动部分在于支持多维度bbox输入，并且shared memory改为动态数组。
  • INT8部分50张图差不多就够了
  • 支持多线程操作

RetinaNet

简介

注意事项

  • 使用转换onnx时候需要设置opset=11
  • 如果在解析onnx时遇到 Assertion failed: ctx->tensors().count(inputName) 这个错误的话，下载最新的onnx-tensorrt源码编译，替换trt对应的lib

ResNet

简介

  • 位置：resnet_main.cpp
  • 对应任意可直接转换的分类模型

FCOS

简介

注意事项

  • 目前trt暂时不支持Group Normalization，如果需要使用GN版本需要单独实现。
  • 有空会更新GN

Hourglass

简介

更新日志

2021.01.26

  1. 增加darknet的trt版本

2020.11.05

  1. 增加了一个KeypointTRT抽象类。截至目前可完成检测、分割、关键点的"一键"模型转换和部署
  2. 实现微软的SimplePose

2020.09.10

  1. 实现StreamProcess，将CPU和GPU端分离实现延迟隐藏，以yolov5和视频流为demon
  2. 线程安全队列可设置成容量有限的队列（避免爆内存）。同时将push和emplace操作改成try_系列，即可能成功或者失败而不是阻塞。
  3. 增加一个nms_cpu的方式

2020.09.03

  1. 重写opencv的bilinear resize算子
  2. 将cvtColor和HWC2CHW融合为一个
  3. 提升cache的命中率
  4. 开启多线程进行图像预处理，Opencv原始的一套 resize + cvtColor 大概15ms。改写之后 resize + cvtColor + HWC2CHW一共4.8ms

2020.08.31

  1. 实现线程安全队列
  2. 实现基于线程安全队列的线程池
  3. 将模型decode部分用线程池代替std::thread实现。应用在yolov3，yolov5，retinaface模型上
  4. 抽空把预处理部分也改了，目前遍历的方式不太能cache shot。

2020.08.26

  1. 整体抽象化，严格按照面向对象将模型剥离出来，提高代码复用率。
  2. Detection类已转换网络：yolov5，yolov3，retinaface，retinanet，fcos
  3. Segmentation类已转换网络：psenet，psev2
  4. 尚未转换的网络：rensnet（分类），hourglass（关键点）

2020.08.19

  1. 重新写了CMAKE，把历史遗留的一些问题解决了。重新组织语言，尽可能用cmake内置的一些变量，同时把依赖的一些路径进行合并，提升工程的可移植性。
  2. 测了OpenMP，在多线程执行任务的速度上不如std::thread，但是需要频繁开启销毁线程时速度比std::thread快。后来查到OpenMP是基于线程池的，故考虑用线程池来代替。

2020.07.13

  1. 增加OCR系列的PANNet，即PSEv2。模型整体轻量化，且不需要像PSE那样设置这么多kernel。 不过讲道理，PSE和PAN在decode的时候都要遍历所有的文本像素点，并没有快很多。测试发现decode部分的实际速度差距已经很小了，感觉FPS提升主要还是换了轻量的backbone。
  2. 这一次将后处理的sigmoid操作没有放到onnx中。同时在网上发现一个fast sigmoid的操作。不过如果走GPU的话差异并不大，走CPU的话速度差了5 6倍。
  3. 改了一下OCR检测输出，既有mask也有RBox

2020.07.04

  1. 增加OCR系列的PSENet。infer时间不算太慢，但是decode部分的渐进式扩张算法耗时太久，这一块其实可以再优化。

2020.06.30

  1. 细节修复，使用cudaEvent来计算时间，而不是cpu

2020.06.17

  1. U版的yolo性能确实好，Python代码封装的也很不错。试了一下yolo5s转成trt速度快了好多，但是准确率也不低
  2. 目前准备转OCR中的ctpn，但是由于代码复现效果一直不好，可能还得等一段时间。
  3. 有空更新一下ncnn的git

2020.05.28

今天开始决定把每次更新的内容记录一下。

  1. 支持多线程（部署在yolo和retinaface上）
  2. 把BRG转RGB和HWC转CHW放在一起操作
  3. worker=4的时候前后处理总用时可以从30ms压缩到15ms左右

About

TensorRT-7 Network Lib 包括常用目标检测、关键点检测、人脸检测、OCR等 可训练自己数据

Topics

hourglass tensorrt retinanet yolov3 psenet fcos retinaface yolov5

Resources

Readme
Activity

Stars

525 stars

Watchers

17 watching

Forks

111 forks
Report repository

Releases

No releases published

Packages

No packages published

Languages