.
├── calibrator.py
├── detect.py
├── doc
│ ├── 3.6res.png
│ ├── ...
├── result
│ ├── mix_res.json
├── export.py
├── LICENSE
├── Logger.py
├── model_save
│ └── yolov5s_calibration.cache
├── README.md
├── tensorrt_engine.py
├── test_image
│ ├── bus.jpg
│ └── zidane.jpg
├── util_trt_modify.py
└── util_trt.py
环境:ubuntu:18.04
cuda:11.1
cudnn:8.0
tensorrt:7.2.3.4
OpenCV:3.4.2
python3.8
Conda
模型包括两部分:
1.strategy
strategy = [4, 4, 4, 4, 8, 4, 4, 4, 4, 4, 4, 8, 4, 4, 8, 4, 4, 8, 4, 4, 8, 4, 4, 4, 8, 8, 4, 4, 4, 8, 4, 8, 4, 4, 4,
4, 8, 4, 6, 4, 8, 4, 4, 8, 6, 4, 4, 8, 4, 8, 8, 6, 4, 4, 8, 4, 4, 8, 4, 8, 4, 4]
2.模型
提取码:wuln
--img_path test_image/bus.jpg
onnx模型保存到pt模型的同一目录下,可以使用netron工具,查看图形化onnx模型
y=model(img)
ort_session = onnxruntime.InferenceSession(f)
ort_outs = ort_session.run(None, ort_inputs)
np.testing.assert_allclose(to_numpy(y[1][0]), ort_outs[0], rtol=1e-03, atol=1e-05)
mse = np.sqrt(np.mean((to_numpy(y[1][0]) - ort_outs[0]) ** 2))
`python export.py --weights model_save/yolov5s.pt --img_path test_image/bus.jpg --img-size 640 --batch-size 1`
tensorrt_engine.py
util_trt_modify.py
calibrator.py
CALIB_IMG_DIR = 'coco/images/train2017'
--onnx_model_path = 'model_save/yolov5s.onnx'
--img_path = 'test_image/bus.jpg'
--engine_model_path = "model_save/yolov5s.trt"
--quantize fp32/fp16/int8/int4
outputs = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
mse = np.sqrt(np.mean((feat[0] - ort_outs[0]) ** 2))
fp32
fp16
int8
mix
python tensorrt_engine.py --quantize fp32 --onnx_model_path model_save/yolov5s.onnx --img_path test_image/bus.jpg --engine_model_path model_save/yolov5s.trt
detect.py
--engine_model_path
--onnx_model_path
--img_path test_image/bus.jpg
yolov5s.trt
yolov5s_fp16.trt
yolov5s_int8.trt
yolov5s_mix.trt
python3 coco_map.py
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.052
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.128
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.036
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.024
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.073
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.051
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.053
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.069
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.070
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.026
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.088
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.076
python detect.py --engine_model_path model_save/yolov5s.trt --img_path test_image/bus.jpg #使用engine检测
python detect.py --onnx_model_path model_save/yolov5s.trt --img_path test_image/bus.jpg #使用onnx检测