前言

把一个 PyTorch 模型导出 ONNX，再用 ATC 编译成 .om，看起来三步就能搞定。实际操作的时候，第三步（ATC 编译）经常报错：「算子不支持」。这个问题在迁移第三方模型（比如 HuggingFace 上的模型）时特别常见。

---

一、哪些 ONNX 算子昇腾不支持

ONNX 的算子规范有 170+ 个，CANN 支持其中约 120 个。不支持的主要是这几类：

1. 稀疏算子

SparseToDense、SparseMatrixMul 这类稀疏计算算子，昇腾目前不支持。原因是达芬奇架构针对稠密计算优化，稀疏计算效率不如 GPU。

替代方案：先转成稠密矩阵，再走标准算子。性能会差一些，但功能等价。

# 原始：稀疏矩阵乘法
output = torch.sparse.mm(sparse_weight, input)

# 替代：转稠密后计算
dense_weight = sparse_weight.to_dense()
output = torch.mm(dense_weight, input)

2. 动态索引算子

GatherND、ScatterND、NonMaxSuppression 这类根据数据内容决定索引位置的算子，部分 shape 下不支持。

替代方案：用 Gather + 预计算索引替代 GatherND。

3. 自定义算子

模型里用了 torch.autograd.Function 或者自定义 C++ 扩展的算子，导出 ONNX 后变成 custom_op，ATC 不认识。

替代方案：用 Ascend C 重写，或者用标准算子组合替代。

常见不支持算子速查表

ONNX 算子	是否支持	替代方案	备注
SparseToDense	❌	转 dense 后用 MatMul	性能损失
GatherND	⚠️ 部分shape	Gather + 预计算索引	shape 固定时可用
NonMaxSuppression	❌	ops-adv 的 NMS 算子	需要单独替换
Resize (cubic)	⚠️ 部分模式	Resize (nearest/linear)	cubic 不支持
HardSigmoid	❌	Clip + Add + Mul	数学等价替换
Swish	❌	Sigmoid + Mul	数学等价替换
TopK (动态 K)	⚠️ K 固定时支持	K 用常量	动态 K 不支持

---

二、算子替换的两种方式

方式一：导出前替换（推荐）

在 PyTorch 里把不兼容的算子替换掉，然后再导出 ONNX。

import torch
import torch.nn as nn

class HardSigmoidReplacement(nn.Module):
    """用 Clip + Add + Mul 替换 HardSigmoid"""
    def forward(self, x):
        # HardSigmoid(x) = clip(x * 1/6 + 0.5, 0, 1)
        return torch.clamp(x * (1.0 / 6.0) + 0.5, 0.0, 1.0)

# 注册替换规则
from torch.onnx import register_custom_op_symbolic

def hardsigmoid_symbolic(g, input):
    # 导出 ONNX 时用 Clip+Add+Mul 代替
    scaled = g.op("Mul", input, g.op("Constant", value_t=torch.tensor(1.0/6.0)))
    shifted = g.op("Add", scaled, g.op("Constant", value_t=torch.tensor(0.5)))
    return g.op("Clip", shifted)

register_custom_op_symbolic("aten::hardsigmoid", hardsigmoid_symbolic, 9)

导出前替换的好处：生成的 ONNX 文件只包含标准算子，ATC 编译零报错。

方式二：导出后修改 ONNX 图

已经导出了 ONNX 文件，不想重新导出。用 onnx-graphsurgeon 修改计算图。

import onnx_graphsurgeon as og

graph = og.import_onnx("model.onnx")

# 找到 HardSigmoid 节点
for node in graph.nodes:
    if node.op_type == "HardSigmoid":
        # 替换成 Clip + Add + Mul
        # ...（图修改逻辑）
        pass

og.export_onnx(graph, "model_fixed.onnx")

这种方式更灵活，但代码量大。适合第三方 ONNX 模型（拿不到 PyTorch 源码的情况）。

---

三、NMS 的特殊处理

目标检测模型几乎都用了 NMS（Non-Maximum Suppression），但 ONNX 的 NonMaxSuppression 算子昇腾不支持。

替换方案

用 ops-adv 仓库的 nms 算子替换。这个算子是昇腾原生实现的，支持 FP16 和 INT8，性能比 CPU 版本快 5 倍。

from ops_adv import nms

# 替换 TorchVision 的 NMS
# 原始：keep = torchvision.ops.nms(boxes, scores, iou_threshold)
# 替换：
keep = nms(boxes, scores, iou_threshold, sorted=True)

导出 ONNX 时，需要注册自定义 symbolic：

def nms_symbolic(g, boxes, scores, iou_threshold):
    return g.op("NMS", boxes, scores, iou_threshold)

register_custom_op_symbolic("ops_adv::nms", nms_symbolic, 11)

性能对比

NMS 实现	1000 个框耗时	精度	设备
TorchVision NMS	8ms	标准	CPU
ONNX NonMaxSuppression	❌ 不支持	-	-
ops-adv NMS	1.5ms	等价	NPU

---

四、动态 shape 模型的 ONNX 导出

动态 shape 模型导出 ONNX 时，要指定动态轴：

torch.onnx.export(
    model,
    dummy_input,
    "model_dynamic.onnx",
    input_names=["input"],
    dynamic_axes={
        "input": {0: "batch", 1: "sequence"},
    },
    opset_version=11
)

ATC 编译时也要指定动态范围：

atc --model=model_dynamic.onnx \
    --framework=5 \
    --output=model_dynamic \
    --input_shape_range="input:[1~8,1~512]"

1~8 表示 batch 在 1 到 8 之间变化，1~512 表示序列长度在 1 到 512 之间变化。ATC 会编译多个子图覆盖这个范围。

注意事项

• 动态范围不要设太宽（比如 1~4096），子图数量爆炸，编译时间不可接受
• 优先把动态维度固定到几个常见值（比如 1,4,8,16），用 --dynamic_batch_size
• 某些算子在动态 shape 下会退化（比如 FlashAttention 要求序列长度是 64 的倍数）

---

五、排查流程

遇到 ATC 编译报错「算子不支持」，按以下流程排查：

1. 查看报错日志，找到不支持的算子名称
   ↓
2. 查 CANN 算子支持列表（$ASCEND_HOME/opp/built-in/op_impl/ai_core/tbe/config/)
   ↓
3a. 支持 → 检查输入 shape 是否满足约束
3b. 不支持 → 查替代方案（数学等价替换 / ops-adv 算子 / Ascend C 自定义）
   ↓
4. 替换后重新导出 ONNX，再编译

---

参考资源

• CANN 算子支持列表：https://www.hiascend.com/document/detail/zh/CANN/
• onnx-graphsurgeon 工具：https://github.com/NVIDIA/TensorRT/tree/main/tools/onnx-graphsurgeon
• ops-adv 算子仓库：https://atomgit.com/cann/ops-adv
• Ascend C 自定义算子开发：https://www.hiascend.com/document/detail/zh/CANN/

---

总结

ONNX 模型迁移到昇腾，算子不兼容是主要障碍。解决思路分三层：能用标准算子数学等价替换的（HardSigmoid、Swish），优先替换；ops-adv 里有原生实现的（NMS），用原生算子；实在没有的，用 Ascend C 自定义。导出前替换比导出后修改 ONNX 图更省事。动态 shape 模型要控制动态范围，太宽会拖慢编译。遇到报错先查算子列表，再找替代方案，不要一上来就写自定义算子。