ATC 做了什么：从 ONNX 到 .om

前言

训练好的模型，怎么跑到昇腾 NPU 上？

答案是 ATC（Ascend Tensor Compiler）。它做的事情很直接：把一个框架导出的模型文件（通常是 ONNX 格式），编译成昇腾 NPU 可以直接执行的 .om 离线模型。

这个编译过程不是简单的"翻译"，而是包含了大量的图级优化。同样的 ONNX 模型，编译出来的 .om 文件，性能可能差几倍------取决于你有没有开优化、用的什么精度策略、有没有做 AOE 调优。

编译流程拆解

一个典型的 ATC 编译过程，经历以下步骤：

第1步：解析模型

读入 ONNX 文件，解析成内存里的一张计算图。这一步会做基本的校验：op type 是否支持、输入输出 shape 是否合法、是否有不支持的算子组合。

第2步：图优化

这是最关键的一步。ATC 会对计算图做多种优化：

算子融合：Conv + BN + ReLU 合成一个 Fused op，减少 kernel 启动次数
常量折叠：初始化就能算出来的子图，直接算好，不生成运行时计算
死代码消除：删掉没有 consumer 的算子节点
数据布局优化：根据达芬奇架构的特点，选择最好的 weight 布局（NC1HWC0 等）

第3步：内存规划

给每个 tensor 分配显存地址。好的内存规划能做到：相邻算子复用同一块显存，峰值显存大幅降低。

第4步：算子选型

给每个算子选择最好的实现。同一个 MatMul，有小 shape 的实现，也有大 shape 的实现；有不用 L2 缓存的，也有专门用 L2 缓存的。ATC 会根据 shape 和硬件配置自动选。

第5步：生成 .om 文件

把所有信息（算子指令、内存布局、调度顺序）打包成 .om 文件，可以直接拿去部署。

从 PyTorch 到 .om 的完整路径

实际操作中最常见的路径是：PyTorch → ONNX → ATC → .om。

Step 1：PyTorch 导出 ONNX

python 复制代码

import torch
import torchvision.models as models

# 1. 加载模型（以 ResNet50 为例）
model = models.resnet50(pretrained=False)
model.eval（)

# 2. 准备 dummy 输入
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出 ONNX
torch.onnx.export(
    model,
    dummy_input,
    "resnet50.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}},  # 支持动态 batch
    opset_version=11
)

print("ONNX 导出完成：resnet50.onnx")

Step 2：用 ATC 编译成 .om

bash 复制代码

# 基础编译命令
atc --model=resnet50.onnx \
    --framework=5 \
    --output=resnet50 \
    --soc_version=Ascend910

# 编译成功会生成 resnet50.om

参数说明：

--framework=5：ONNX 对应的 framework ID
--soc_version=Ascend910：目标芯片型号
--output：输出的 .om 文件名（不用加后缀）

Step 3：验证 .om 文件

bash 复制代码

# 用 atc 自带的模型查看工具
omg_info resnet50.om

# 预期输出（示例）：
# Model Name: resnet50
# Input  : input [1, 3, 224, 224] float32
# Output : output [1, 1000] float32
# Op Num : 53

进阶：开 AOE 调优再编译

基础编译出来的 .om 能用，但不一定最快。要做性能优化，需要先用 AOE（Ascend Optimization Engine）做自动调优，再把调优结果喂给 ATC。

bash 复制代码

# Step 1: 用 AOE 做算子级调优
atc --model=resnet50.onnx \
    --framework=5 \
    --output=resnet50_aoe \
    --soc_version=Ascend910 \
    --auto_tune_mode="GA,RL" \
    --tuning_iterations=50

# 调优过程可能需要 30 分钟到数小时
# 调优结果会缓存到 ~/.ascend/aoe/

# Step 2: 用调优结果重新编译
atc --model=resnet50.onnx \
    --framework=5 \
    --output=resnet50_optimized \
    --soc_version=Ascend910 \
    --auto_tune_mode="GA,RL" \
    --load_tuning_result=~/.ascend/aoe/

调优后的 .om 文件，推理延迟通常能降低 10-30%。

精度策略：什么时候用 FP16

ATC 编译时可以指定精度策略：

bash 复制代码

# 允许把 FP32 算子自动转成 FP16（提速，但可能有精度损失）
atc --model=model.onnx \
    --framework=5 \
    --output=model_fp16 \
    --soc_version=Ascend910 \
    --op_precision_mode=allow_fp32_to_fp16

这个开关对大模型特别有用。LLaMA 推理时，90% 以上的算子都可以用 FP16 算，只有少部分（Softmax、LayerNorm）需要保留 FP32 保精度。

op_precision_mode 的可选值：

force_fp16：强制全部转 FP16（速度快，精度风险高）
allow_fp32_to_fp16：自动判断，能转则转
keep_origin：保持原精度（最安全，但可能慢）

常见编译错误与排查

错误1：unsupported op type:xxx

原因：ONNX 模型里用到了 ATC 还没支持的算子。

解决：

bash 复制代码

# 查看 ATC 支持的算子清单
atc --help_op

# 如果确实不支持，有两个选择：
# 1. 在 PyTorch 里把这个算子拆成多个已支持的算子
# 2. 用 Ascend C 自定义算子，注册到 ATC 里

错误2：shape inference failed

原因：ONNX 导出时某些算子的输出 shape 推导失败，ATC 无法做内存规划。

解决：在 torch.onnx.export 时加上 dynamic_axes 参数，或者把模型改成静态 shape 再导出。

错误3：编译出来的 .om 推理精度不对

原因：通常是算子融合策略导致的数值偏差。

解决：

bash 复制代码

# 关闭算子融合，重新编译
atc --model=model.onnx \
    --framework=5 \
    --output=model_no_fuse \
    --soc_version=Ascend910 \
    --enable_fusion=False

关掉融合后如果精度恢复正常，说明是某个融合算子实现有 bug，可以提 issue 给 cann/atc 仓库。

总结

ATC 编译器是把 PyTorch/ONNX 模型部署到昇腾 NPU 的必经之路。它的核心价值不在于"翻译"，而在于编译期的图优化------通过算子融合、内存规划、算子选型等手段，让最终生成的 .om 模型在 NPU 上跑得更快。

实际项目中，ATC 编译通常跟 AOE 调优配合使用：先让 AOE 找到最优的算子参数，再把调优结果喂给 ATC 生成最终模型。这个流程多花几小时，但换来的推理性能提升是值得的。