模型部署源代码(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 :保存权重文件,请下载附件获取模型权重文件
该仓库主要代码分为两大部分
文件名: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:转换好的模型序列化保存需要用到的名字,建议将输入图像尺寸作为名字的一部分。
输出:转换、序列化后的模型文件。
具体流程可以参考下图
矩形推理发生的地方位于两个阶段,代码地址分别是:utils.datasets.py和models_rect.py中:
dataset = LoadImages(path=opt.source, img_size=(opt.img_size, opt.img_size), auto=opt.rect)
yolo_head = YOLOHead(config_path=opt.model_def)
设备型号 | 运行模型 | 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 |