【Debug】YOLOv8模型TensorRT INT8 量化踩坑

环境：

win10 + anaconda 虚拟环境；
tensorrt：8.6.1；
onnx：1.15.0；
onnxslim：0.1.31；
CUDA Version：12.2；
torch：2.1.0；
ultralytics：v8.2.45；

---

问题描述

YOLOv8项目下导出自己改进的模型【TensorRT INT8 格式】报错

在YOLOv8项目下导出TensorRT格式的模型

from ultralytics import YOLO
# Load the YOLOv8 model
model = YOLO("yolov8n-obb.pt")
# TensorRT INT8 with calibration `data` (i.e. COCO, ImageNet, or DOTAv1 for appropriate model task)
out = model.export(
    format="engine", 
    imgsz=640, 
    dynamic=True, 
    verbose=False, 
    batch=8, 
    workspace=0, 
    int8=True, 
    data="dota8.yaml"
)

使用官方给出的旋转目标检测模型yolov8n-obb，导出与量化没有问题，一切顺利。
！！但是！！替换成自己改进后的模型，导出过程出现如下报错

[06/28/2024-17:04:08] [TRT] [V] Calculating Maxima
[06/28/2024-17:04:08] [TRT] [I] Starting Calibration.
[06/28/2024-17:04:09] [TRT] [E] 1: [softMaxV2Runner.cpp::nvinfer1::rt::task::CaskSoftMaxV2Runner::execute::226] Error Code 1: Cask (shader run failed)
[06/28/2024-17:04:09] [TRT] [E] 3: [engine.cpp::nvinfer1::rt::Engine::~Engine::298] Error Code 3: API Usage Error (Parameter check failed at: engine.cpp::nvinfer1::rt::Engine::~Engine::298, condition: mExecutionContextCounter.use_count() == 1. Destroying an engine object before destroying the IExecutionContext objects it created leads to undefined behavior.
)
[06/28/2024-17:04:09] [TRT] [E] 2: [calibrator.cpp::nvinfer1::builder::calibrateEngine::1181] Error Code 2: Internal Error (Assertion context->executeV2(&bindings[0]) failed. )
TensorRT: export failure ❌ 30.2s: __enter__
Traceback (most recent call last):
 ···
    with build(network, config) as engine, open(f, "wb") as t:
AttributeError: __enter__

进程已结束，退出代码为 1

一直卡在Calibration这里，耽误了我半天时间。

开始以为是
with build(network, config) as engine, open(f, "wb") as t:
这部分代码问题，查了半天，没有找到好的解决方法。【是我太菜了】

然后试了不做INT8量化的导出，居然可以正常导出？？
但是推理的使用又出现下面的报错。。。
检测不出来任何东西。

[TRT] [E] 1: [softMaxV2Runner.cpp::nvinfer1::rt::task::CaskSoftMaxV2Runner::execute::226] Error Code 1: Cask (shader run failed)

突然想到tensorrt可以自定义算子，那是不是我的改进部分的代码，tensorrt原来是不支持的？
我需要改成tensorrt兼容的格式？

---

原因分析：

因为对模型的检测头做了改进，所很快定位到应该是网络模型的改动，【也没很快，定位定了半天】
导致有部分算子/计算流程没有很好的导出，使得模型与量化流程不兼容。

---

解决方案：

定位了一下检测头中改进的部分，
发现一些计算响上述的量化流程。
啪的一下就找到，很快嗷，就发现下面的代码会影响模型导出

idx = torch.cos(angle_a) * torch.cos(angle_b) + torch.sin(angle_a) * torch.sin(angle_b) < 0

这行代码又是三角函数，又是条件判断的！！！
统统拆开！！！【谢谢GPT，都不用我自拆了】
因为在 TensorRT 中，某些操作在 INT8 模式下的支持可能有限或需要特殊处理。
把上述的计算拆解，最后就成功导出INT8量化格式的模型啦~~

TensorRT: export success ✅ 287.7s, saved as '...\best.engine' (6.3 MB)

Export complete (288.7s)
Results saved to ...
Predict:         yolo predict task=obb model=...\best.engine imgsz=640 int8 
Validate:        yolo val task=obb model=...\best.engine imgsz=640 data=...\data.yaml int8 
Visualize:       https://netron.app

参考
https://github.com/ultralytics/ultralytics