/YOLOV3_Tensorrt_Rectangular_Inference

After use TensorRT model & rectangular Inference , the time forward is much lower then before.

Primary LanguagePythonMIT LicenseMIT

模型部署源代码（Jetson Nano）说明

[toc]

项目结构图

模型部署源代码（Jetson Nano）

│ detect_demo.py ：演示demo

│ models.py ：只支持方形推理的模型

│ models_rect.py ：支持方形推理和矩形推理的模型

│ README.md ：说明文档

│ torch_model_2_trt.py ：模型转换demo

│

├─config

│ 请下载附件获取模型结构文件（cfg）

│

├─data

│ │ helmet.names ：类别，hat,person

│ │

│ ├─custom

│ │ ├─images ：保存图片

│ │ └─labels ：保存标签

│ ├─output ：保存检测结果

│ └─samples ：保存测试原数据

├─utils

│ datasets.py ：进行数据加载

│ parse_config.py ：解析cfg

│ utils.py ：工具包

└─weights ：保存权重文件,请下载附件获取模型权重文件

代码使用说明

  • 1.准备好YOLOv3模型以及相应cfg
  • 2.执行 torch_model_2_trt.py，模型转换，注意，转换TensorRT的时候需要设置X的形状，如：[1,3,416,256]，后两位表示输入形状。
  • 3.执行detect_demo.py ，进行模型推理，注意，输入图片形状必须是TensorRT对应的形状，否则TensorRT模型的输出是无效的。

代码文件说明

  • 该仓库主要代码分为两大部分
    • 1.模型转换Demo
    • 2.模型推理Demo

模型转换Demo

  • 文件名：torch_model_2_trt.py

  • 概述：该文件功能是将pytorch的模型转换成包含TensorRT加速引擎的TRTModule模型，在GPU层面上实现加速效果。

  • 主要参数说明：

    • model_def：模型config文件地址

    • weights_path：权重地址，权重可能有两种后缀，分别为.weights和.pth，无论那种，cfg必须是对应的

    • half：半精度

    • device：Jetson Nano拥有一块128core的MaxwellGPU，因此选GPU时不需要选块号，以T/F表示GPU/CPU

    • img_size:模型需要定义输入形状，默认是正方形，如果需要转换成矩形，请输入相应的形状，由于YOLO特征层的特点，输入形状的每个值都必须是32的倍数。

      注意：如果输入形状不对应，检测会失败

    • model_save_name：转换好的模型序列化保存需要用到的名字，建议将输入图像尺寸作为名字的一部分。

  • 输出：转换、序列化后的模型文件。

模型演示Demo

  • 文件名：detect_demo.py
  • 概述：该文件功能是演示在Nano上从：读取视频——>推理——>输出的整个过程。主体**是：利用CPU和GPU独立运行的特性使得三大步骤并行执行
  • 主要参数说明：
    • model_def、weights_path、half参考上文
    • rect：控制矩形推理是否执行
    • tensorrt：控制是否使用TensorRT模型
    • conf_thres：置信度阈值
    • iou_thres：iou阈值
    • batch_size：每个batch的大小，默认1
    • img_size：长边大小416，若输入是不是正方形，则长边为416，短边为等比例缩小后最小的32的倍数
    • view_img：是否输出检测结果图像
  • 输出：控制台输出帧率（视频左上角也有），检测结果视频播放

部分加速原理说明

关于本项目矩形推理加速的原理说明

  • 具体流程可以参考下图

  • 来自博客的图

  • 如想看具体过程，可参考我的博客：椰子奶糖的Jetson nano学习历程

  • 矩形推理发生的地方位于两个阶段,代码地址分别是:utils.datasets.py和models_rect.py中：

    • 1、LoadImages：dataset = LoadImages(path=opt.source, img_size=(opt.img_size, opt.img_size), auto=opt.rect)
    • 2、YOLOLayer的forward：yolo_head = YOLOHead(config_path=opt.model_def)
      • 本项目预留了原版models，不支持矩形推理，可与新版models_rect.py对比查看

本项目数据预加载加速原理说明

  • 利用CPU与GPU可以并行的特点，将下一帧的数据加载与本帧模型推理放在同一时间处理，此时耗时较短的数据加载便会在宏观上消失，具体参考detect_demo.py中的while vid_cap部分代码。
    • 更新：后处理NMS也加入进来，目前是:（上一帧的NMS+下一帧的数据加载）与（本帧的模型推理）并行。

Jetson Nano上效果对比

  • Nano的算力并不算强，Maxwell这块U算力大约是1060的$$\frac{1}{7} - \frac{1}{10}$$
  • 目前做过如下测试：
设备型号 运行模型 GPU TensorRT加速 剪枝 矩形推理 FPS
GTX1060 YOLOV3-SPP √ 18.4
Jetson Nano YOLOV3-SPP 0.3
Jetson Nano YOLOV3-SPP √ 2
Jetson Nano YOLOV3-SPP √ √ 3.7
Jetson Nano YOLOV3-SPP √ √ √ 12.3
Jetson Nano YOLOV3-SPP √ √ √ √ 19.2