### 解决 YOLOv8 导出 ONNX 模型时生成 INT64 类型的问题 当导出 YOLOv8 模型为 ONNX 格式时，可能会遇到权重被定义为 `INT64` 的情况。然而，TensorRT 并不原生支持 `INT64` 数据类型，在这种情况下会尝试将其转换为 `INT32`[^1]。 为了确保模型能够顺利部署并兼容 TensorRT 推理引擎，可以采取以下措施来处理这个问题： #### 方法一：修改源码以强制使用 INT32 可以在 PyTorch 中通过调整网络结构中的张量数据类型来避免生成 `INT64` 权重。具体来说，可以通过遍历整个模型图并将所有整数类型的参数设置为 `torch.int32` 或者在保存前显式指定输出节点的数据类型为 `int32`。 ```python import torch from ultralytics import YOLO def convert_model_to_int32(model): for param in model.parameters(): if isinstance(param.data, (torch.IntTensor, torch.LongTensor)): param.data = param.data.to(torch.int32) model = YOLO("runs/detect/train/weights/best.pt") convert_model_to_int32(model) model.export(format="onnx", opset_version=12) # 使用适当版本号 ``` #### 方法二：利用 ONNX 工具链优化 另一种方式是在完成初步导出之后，借助于 ONNX 提供的各种工具来进行后续的优化操作。比如使用 `onnxmltools` 库或者其他第三方库对已有的 `.onnx` 文件做进一步加工，从而达到改变特定层或变量属性的目的。 ```bash pip install onnxmltools ``` ```python from onnxmltools.utils.float16_converter import convert_float_to_float16 from onnx import load_model, save_model # Load the original ONNX model. original_model_path = 'path/to/original/model.onnx' optimized_model_path = 'path/to/save/optimized_model.onnx' onnx_model = load_model(original_model_path) # Convert all float tensors to float16 and int64 tensors to int32 within this function call. converted_model = convert_float_to_float16(onnx_model, keep_io_types=True, min_positive_val=1e-7, max_finite_val=1e4) save_model(converted_model, optimized_model_path) ``` 上述两种方案都可以有效地帮助解决由于 `INT64` 而引起的潜在问题。如果希望更深入地控制导出过程，则建议直接编辑框架内部的相关部分；而对于大多数应用场景而言，第二种基于现有工具链的方式可能更为便捷高效[^2][^3].