前言
训练好的模型,怎么跑到昇腾 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
importtorchimporttorchvision.modelsasmodels# 1. 加载模型(以 ResNet50 为例)model=models.resnet50(pretrained=False)model.eval()# 2. 准备 dummy 输入dummy_input=torch.randn(1,3,224,224)# 3. 导出 ONNXtorch.onnx.export(model,dummy_input,"resnet50.onnx",input_names=["input"],output_names=["output"],dynamic_axes={"input":{0:"batch_size"}},# 支持动态 batchopset_version=11)print("ONNX 导出完成:resnet50.onnx")Step 2:用 ATC 编译成 .om
# 基础编译命令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 文件
# 用 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。
# 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 编译时可以指定精度策略:
# 允许把 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 还没支持的算子。
解决:
# 查看 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 推理精度不对
原因:通常是算子融合策略导致的数值偏差。
解决:
# 关闭算子融合,重新编译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 生成最终模型。这个流程多花几小时,但换来的推理性能提升是值得的。