yolov5模型转换为TensorRT模型时报错Your ONNX model has been generated with INT64 weights

1. 报错信息

今天在yolov5模型转换为TensorRT模型时报错：

Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64

详细报错信息如下：

----------------------------------------------------------------
Input filename:   ../models/yolov5s.onnx
ONNX IR version:  0.0.6
Opset version:    13
Producer name:    pytorch
Producer version: 1.9
Domain:           
Model version:    0
Doc string:       
----------------------------------------------------------------
[2023-03-08 07:33:55 WARNING] onnx2trt_utils.cpp:198: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.
[2023-03-08 07:33:55 WARNING] onnx2trt_utils.cpp:198: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.
[2023-03-08 07:33:55 WARNING] onnx2trt_utils.cpp:198: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.

 

2. 原因分析

根据提示，我们可以发现原因：

因为我们的onnx模型参数类型是INT64, 然后TensorRT是不支持 INT64的。这就需要我们先把ONNX模型的精度降到 INT32，然后再进行转换为 TensorRT模型。

3. 解决办法

具体操作如下：

先对onnx模型进行onnxsim 简化：

python -m onnxsim yolov5s.onnx yolov5s_sim.onnx

./trtexec --onnx=yolov5s_sim.onnx --saveEngine=yolov5s.engine

这样就解决了