人在途中：从“编译失败”到“模型可落地”——CUDA 自定义算子

1. 问题现象：模型能跑，但落不了地

在项目支持过程中，一个非常典型但高频的问题是：

模型在 PC 侧验证一切正常，但在地平线工具链编译阶段直接失败。

典型表现：

• ONNXRuntime 推理 ✅ 正常
• TensorRT 推理 ✅ 正常
• hb_compile ❌ 失败（算子不支持）

进一步分析后，基本都能定位到一个共性问题：

👉 ONNX 模型中存在 自定义算子（Custom Operator）

2.不要误判：这不是工具链问题

很多客户第一反应是：

• 工具链版本是不是太老？
• 参数是不是配错了？
• 能不能加 plugin 支持？

这里需要明确一个结论：

❗ 这是“模型表达问题”，不是“工具链能力问题”

3.本质拆解：为什么会失败？

3.1 ​ONNX​​ 并不等价

ONNX 只是一个中间表示，但​表达能力是分层的​：

类型	描述	工具链支持
标准算子 ONNX	可解析计算图	✅
自定义算子 ONNX	黑盒节点	❌

3.2 Runtime vs Compiler

系统	本质	能力
ONNXRuntime	解释执行	支持自定义 kernel
TensorRT	插件扩展	支持 plugin
地平线工具链	静态编译	❌ 不支持自定义算子
👉 核心差异：

Runtime 可以“执行”，Compiler 必须“理解”

4.torch​ 化：不是建议，是必选项

4.1 ​torch​​ 化定义

用标准 ​​PyTorch​​ 算子组合，重写自定义算子的计算逻辑

4.2 从计算图角度看

原始：

CustomOp(x)

torch 化：

x → add → mul → reshape → ...

4.3 本质三件事

1. 语义展开（把黑盒拆开）
2. 算子对齐（只用标准算子）
3. 数值等价（结果一致）

5. 示例说明（核心）

5.1 原始模型

class ModelWithCuda(nn.Module):
    def forward(self, x, y):
        return custom_add(x, y)

特点：

• custom_add = CUDA kernel
• ONNX 中为自定义算子

5.2 torch 化模型

class ModelTorch(nn.Module):
    def forward(self, x, y):
        return x + y

5.3 实验结果

项目	CUDA ONNX	torch ONNX
ONNXRuntime	✅	✅
工具链编译	❌	✅
数值一致性	✔	✔
DEMO 示例：horizon_cuda_torchify.tar.gz

6.工程级 torch 化流程（关键部分）

Step 1：检测自定义算子

python3 check_custom_ops.py model.onnx

Step 2：理解算子逻辑

来源可能是：

• CUDA kernel
• TRT plugin
• PyTorch Extension

必须搞清：

• 输入
• 输出
• 中间计算

Step 3：设计替代实现

类型	替代方式
简单算子	直接替换
复合算子	拆分实现
特殊算子	数学等价变换

Step 4：重写模型

class TorchImpl(nn.Module):
    def forward(self, x):
        return torch_impl(x)

Step 5：导出 ONNX

python3 export_torch_onnx.py

Step 6：验证（必须）

• ONNXRuntime 推理
• 数值一致性（MSE / cosine）
• hb_compile 编译

7. 常见错误（非常关键）

❌ 只替换部分算子

→ ONNX 仍包含 custom op

❌ 使用不支持算子

→ 编译仍失败

❌ 忽略数值一致性

→ 精度问题

❌ 误以为导出成功就算完成

→ 必须通过编译

8. 完成标准

必须同时满足：

• ONNX 无自定义算子
• ONNXRuntime 正常
• hb_compile 成功

9. 实战建议

9.1 模型设计阶段

避免：

• CUDA Extension
• TRT Plugin

9.2 项目流程中加入检查

增加：

• ONNX 合法性检测
• 自动算子扫描

10.一句话总结

torch 化不是优化手段，而是模型落地的前提条件

11.结语

这个问题的本质不是“工具链不支持”，而是：

模型没有用硬件可以理解的方式表达