/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