ATC 做了什么:从 ONNX 到 .om
前言训练好的模型怎么跑到昇腾 NPU 上答案是 ATCAscend 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 1PyTorch 导出 ONNXimporttorchimporttorchvision.modelsasmodels# 1. 加载模型以 ResNet50 为例modelmodels.resnet50(pretrainedFalse)model.eval)# 2. 准备 dummy 输入dummy_inputtorch.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_version11)print(ONNX 导出完成resnet50.onnx)Step 2用 ATC 编译成 .om# 基础编译命令atc--modelresnet50.onnx\--framework5\--outputresnet50\--soc_versionAscend910# 编译成功会生成 resnet50.om参数说明--framework5ONNX 对应的 framework ID--soc_versionAscend910目标芯片型号--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能用但不一定最快。要做性能优化需要先用 AOEAscend Optimization Engine做自动调优再把调优结果喂给 ATC。# Step 1: 用 AOE 做算子级调优atc--modelresnet50.onnx\--framework5\--outputresnet50_aoe\--soc_versionAscend910\--auto_tune_modeGA,RL\--tuning_iterations50# 调优过程可能需要 30 分钟到数小时# 调优结果会缓存到 ~/.ascend/aoe/# Step 2: 用调优结果重新编译atc--modelresnet50.onnx\--framework5\--outputresnet50_optimized\--soc_versionAscend910\--auto_tune_modeGA,RL\--load_tuning_result~/.ascend/aoe/调优后的.om文件推理延迟通常能降低 10-30%。精度策略什么时候用 FP16ATC 编译时可以指定精度策略# 允许把 FP32 算子自动转成 FP16提速但可能有精度损失atc--modelmodel.onnx\--framework5\--outputmodel_fp16\--soc_versionAscend910\--op_precision_modeallow_fp32_to_fp16这个开关对大模型特别有用。LLaMA 推理时90% 以上的算子都可以用 FP16 算只有少部分Softmax、LayerNorm需要保留 FP32 保精度。op_precision_mode的可选值force_fp16强制全部转 FP16速度快精度风险高allow_fp32_to_fp16自动判断能转则转keep_origin保持原精度最安全但可能慢常见编译错误与排查错误1unsupported op type:xxx原因ONNX 模型里用到了 ATC 还没支持的算子。解决# 查看 ATC 支持的算子清单atc--help_op# 如果确实不支持有两个选择# 1. 在 PyTorch 里把这个算子拆成多个已支持的算子# 2. 用 Ascend C 自定义算子注册到 ATC 里错误2shape inference failed原因ONNX 导出时某些算子的输出 shape 推导失败ATC 无法做内存规划。解决在torch.onnx.export时加上dynamic_axes参数或者把模型改成静态 shape 再导出。错误3编译出来的 .om 推理精度不对原因通常是算子融合策略导致的数值偏差。解决# 关闭算子融合重新编译atc--modelmodel.onnx\--framework5\--outputmodel_no_fuse\--soc_versionAscend910\--enable_fusionFalse关掉融合后如果精度恢复正常说明是某个融合算子实现有 bug可以提 issue 给cann/atc仓库。总结ATC 编译器是把 PyTorch/ONNX 模型部署到昇腾 NPU 的必经之路。它的核心价值不在于翻译而在于编译期的图优化——通过算子融合、内存规划、算子选型等手段让最终生成的.om模型在 NPU 上跑得更快。实际项目中ATC 编译通常跟 AOE 调优配合使用先让 AOE 找到最优的算子参数再把调优结果喂给 ATC 生成最终模型。这个流程多花几小时但换来的推理性能提升是值得的。