# tensorrt-plugin

**Repository Path**: yongzi-huang/tensorrt-plugin

## Basic Information

- **Project Name**: tensorrt-plugin
- **Description**: TensorRT 自定义算子实验：图手术 + Plugin 开发 | TensorRT / ONNX
· 用 onnx_graphsurgeon 做图手术，将 YOLO11n 中 77 处 Sigmoid+Mul 融合为自定义算子 SigmoidMul_TRT
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-09-25
- **Last Updated**: 2026-10-02

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# TensorRT 自定义算子实验：图手术 + 自定义 Plugin 开发

基于开源项目 [TensorRT-YOLO](https://github.com/laugh12321/TensorRT-YOLO) 的部署实验。

**核心工作**：把 YOLO11n head 里的 77 处 `Sigmoid + Mul` 融合成自定义算子 `SigmoidMul_TRT`，
并把 TensorRT 10.x 弃用的 `EfficientNMS_TRT` 插件迁移到标准 `NonMaxSuppression`。

> **说明**：本仓库只包含我写的实验代码（图手术脚本、自定义 Plugin、构建/推理脚本）。
> TensorRT-YOLO 主体代码（编译环境、官方示例、CMake 顶层）不在此仓库，请从原仓库获取。
> 因此本仓库**不能独立运行**，依赖 TensorRT-YOLO 完整工程。

---

## 一、核心数据

### 图手术前后对比

| op 类型 | 手术前 | 手术后 |
|---|---|---|
| Sigmoid | 78 | 1 |
| Mul | 79 | 2 |
| **SigmoidMul_TRT** | 0 | **77** |
| **NonMaxSuppression** | 0 | **1** |
| EfficientNMS_TRT | 1 | 0 |
| **节点总数** | **322** | **323** |

### 推理结果（YOLO11n，image/output_bus.jpg）

```

检测到 5 个目标:
bus:    0.939   box=(12, 228, 799, 735)
person: 0.902   box=(49, 398, 243, 904)
person: 0.849   box=(671, 393, 810, 880)
person: 0.833   box=(223, 406, 345, 860)
person: 0.396   box=(0, 550, 66, 872)

```

### Engine 构建

| 指标 | 数值 |
|---|---|
| 网络层数 | 415 |
| 构建耗时 | 268 秒 |
| 输入 | `images [1, 3, 640, 640]` |
| 输出 | `selected_indices [-1, 3]` + `boxes [1, 8400, 4]` + `scores [1, 8400, 80]` |

---

## 二、完整流水线

```

yolo11n.pt（Ultralytics 权重）
│
│  yolo export + trtyolo-export（外部工具）
▼
yolo11n-trtyolo.onnx（含 EfficientNMS_TRT）
│
│  graph_surgery/fuse_sigmoid_mul.py（本仓库）
▼
yolo11n-sigmoidmul.onnx（含 EfficientNMS_TRT + 77× SigmoidMul_TRT）
│
│  graph_surgery/replace.py（本仓库）
▼
yolo11n-fused.onnx（含 NonMaxSuppression + 77× SigmoidMul_TRT）
│
│  apps/build_engine.py（本仓库）
▼
yolo11n_fused.engine
│
│  apps/detect1.py（本仓库）
▼
output_bus.jpg（带检测框）

```

**两次图手术（fuse + replace）+ 一次 Engine 构建 + 一次推理**，全部在本仓库。

---

## 三、关键技术点

### Step 1：融合 77 处 Sigmoid + Mul

**做了什么**：用 `onnx_graphsurgeon` 扫描 ONNX 图，找到所有 `Mul` 节点，
把它的输入追溯到 `Sigmoid` 节点，合并成一个 `SigmoidMul_TRT` 自定义节点。

**核心逻辑**：

```python
fused = gs.Node(
    op="SigmoidMul_TRT",
    inputs=[sig_node.inputs[0], other_input],
    outputs=[node.outputs[0]],
    attrs={"plugin_version": "1", "plugin_namespace": ""},
)
# 关键：断开旧节点的连接再删除
node.outputs.clear()
sig_node.outputs.clear()
graph.nodes.remove(node)
graph.nodes.remove(sig_node)
```

遇到的坑：如果只 graph.nodes.remove(node) 而不先 node.outputs.clear()，
graph.cleanup() 会报 Encountered a node not in the graph。



### Step 2：NMS 插件迁移（EfficientNMS_TRT → NonMaxSuppression）

动机：TensorRT 10.x 弃用了 EfficientNMS_TRT 插件。构建 Engine 时报：

```
Could not find any implementation for node EfficientNMS_TRT
```

解法：替换成 ONNX 标准算子 NonMaxSuppression，TensorRT 会自动转成 INMSLayer。

两个关键：

1：输入布局不一致

算子 scores 布局
EfficientNMS_TRT [B, N_boxes, N_classes]
NonMaxSuppression [B, N_classes, N_boxes]

解决：在 scores 前插入 Transpose(perm=[0, 2, 1])。

2：输出格式不同

算子 输出
EfficientNMS_TRT 4 个张量（num_dets, boxes, scores, classes）
NonMaxSuppression 1 个张量 [N, 3]（索引）

解决：把 boxes 和 scores 也暴露成 Engine 输出，Python 侧用索引 Gather 出结果。


### Step 3：从零实现 IPluginV3

目录：sigmoidMulPlugin/

| 文件 | 作用 |
|---|---|
| `SigmoidMulPlugin.h / .cpp` | Plugin 类，实现 `IPluginV3` + `OneCore` / `OneBuild` / `OneRuntime` 三个能力接口 |
| `SigmoidMulPluginCreator.h / .cpp` | Creator，实现 `createPlugin` 和 `getFieldNames` |
| `SigmoidMulKernel.cuh / .cu` | CUDA kernel：`out[i] = sigmoid(sig_in[i]) * w[i]`，含空指针和 `n<=0` 边界保护 |
| `PluginRegistration.cpp` | `REGISTER_TENSORRT_PLUGIN(SigmoidMulPluginCreator)` 宏，触发静态注册 |
| `CMakeLists.txt` | 调用项目自带的 `configure_cuda_trt_target` 完成链接 |

### Step 4：变异测试验证 Plugin 真的在跑

方法：把 kernel 里的 SiLU 实现故意改成 out[i] = sig_in[i] * 3.0f，
重编 Plugin（不重建 Engine），重新推理。

结果：检测置信度全部变成 NaN，框全乱。

结论：kernel 确实在被 TensorRT 调用。改回正确实现后推理恢复正常。

---


## 四、文件结构

```
tensorrt-plugin/
├── README.md
├── .gitignore
├── image
    ├── bus.jpg                    # 原图
    ├── output_bus.jpg             # 结果
├── graph_surgery/
│   ├── fuse_sigmoid_mul.py        # 融合 77 处 Sigmoid + Mul
│   └── replace.py                 # EfficientNMS_TRT → NonMaxSuppression
sigmoidMulPlugin/
├── SigmoidMulPlugin.h          # Plugin 类声明
├── SigmoidMulPlugin.cpp        # Plugin 类实现
├── SigmoidMulPluginCreator.h   # Creator 声明
├── SigmoidMulPluginCreator.cpp # Creator 实现
├── SigmoidMulKernel.cuh        # CUDA kernel 声明
├── SigmoidMulKernel.cu         # CUDA kernel 实现
├── PluginRegistration.cpp      # REGISTER_TENSORRT_PLUGIN 注册
└── CMakeLists.txt              # 编译配置
└── apps/
    ├── build_engine.py            # 用自编译 Plugin 构建 TensorRT Engine
    └── detect1.py                 # 推理 + Python 侧解码
```


---

## 六、环境依赖

本仓库不能独立运行，依赖以下环境：

1. TensorRT-YOLO 完整工程：https://github.com/laugh12321/TensorRT-YOLO
2. TensorRT 10.1：自编译 libnvinfer_plugin.so
3. CUDA 12.2
4. Python 依赖：tensorrt、pycuda、onnx、onnx_graphsurgeon、opencv-python

集成步骤

```bash
# 1. Clone TensorRT-YOLO 并编译
git clone https://github.com/laugh12321/TensorRT-YOLO
cd TensorRT-YOLO
cmake -B build -DTRT_PATH=/path/to/tensorrt -DBUILD_PYTHON=ON
cmake --build build -j4

# 2. 把本仓库的 Plugin 复制到 TensorRT-YOLO
cp -r sigmoidMulPlugin /path/to/TensorRT-YOLO/modules/plugin/

# 3. 在 modules/plugin/CMakeLists.txt 里加：
#    add_subdirectory(sigmoidMulPlugin)

# 4. 重新编译
cmake --build build -j4
```

运行流水线

```bash
# 前置：获取 ONNX（外部工具，非本仓库）
yolo export model=models/yolo11n.pt format=onnx batch=1
trtyolo-export -i models/yolo11n.onnx -o models/yolo11n-trtyolo.onnx -s

# 本仓库的两次图手术
python graph_surgery/fuse_sigmoid_mul.py
python graph_surgery/replace.py

# 构建 Engine
python apps/build_engine.py

# 推理
python apps/detect1.py
```