拿到一个 ONNX 文件，运行 ATC 编译成 .om，几十秒后拿到可执行文件。这个过程里 ATC 做了什么？下面用代码和命令拆解每个步骤。

---

一、ATC 编译命令

最基本的编译命令：

atc --model=model.onnx \
    --framework=5 \
    --output=model \
    --input_shape="input:1,3,224,224"

参数说明：

• --framework=5：ONNX 框架
• --output：输出文件名（自动加 .om 后缀）
• --input_shape：输入 tensor 的 shape

常用编译参数

# 完整编译参数示例
atc --model=resnet50.onnx \
    --framework=5 \
    --output=resnet50 \
    --input_shape="input:1,3,224,224" \
    --op_precision_mode=allow_fp32_to_fp16 \
    --enable_fusion=true \
    --log=info \
    --dump=1

---

二、第一步：图解析

ATC 读入 ONNX 文件，把 proto 格式转成 GE 内部的 ComputeGraph。

查看 ONNX 模型结构

import onnx

# 加载 ONNX 模型
model = onnx.load("resnet50.onnx")

# 打印模型信息
print(f"IR Version: {model.ir_version}")
print(f"Opset Version: {model.opset_import[0].version}")
print(f"Inputs: {[i.name for i in model.graph.input]}")
print(f"Outputs: {[o.name for o in model.graph.output]}")
print(f"Nodes: {len(model.graph.node)}")

输出示例：

IR Version: 7
Opset Version: 11
Inputs: ['input']
Outputs: ['output']
Nodes: 178

ONNX 节点遍历

# 遍历所有算子节点
for node in model.graph.node:
    print(f"Op: {node.op_type}, Name: {node.name}")
    print(f"  Inputs: {list(node.input)}")
    print(f"  Outputs: {list(node.output)}")

输出示例：

Op: Conv, Name: conv1
  Inputs: ['input', 'conv1.weight', 'conv1.bias']
  Outputs: ['conv1.output']
Op: BatchNormalization, Name: bn1
  Inputs: ['conv1.output', 'bn1.weight', 'bn1.bias', 'bn1.running_mean', 'bn1.running_var']
  Outputs: ['bn1.output']
Op: Relu, Name: relu1
  Inputs: ['bn1.output']
  Outputs: ['relu1.output']

算子映射表

ATC 会把 ONNX 算子映射到 CANN 算子：

# ONNX 到 CANN 的算子映射（部分）
OP_MAPPING = {
    "Conv": "Convolution",
    "BatchNormalization": "BatchNorm",
    "Relu": "ReLU",
    "MaxPool": "Pooling",
    "Gemm": "MatMul",
    "Add": "Add",
    "Mul": "Mul",
    "Softmax": "Softmax",
    "MatMul": "MatMul",
    "Reshape": "Reshape",
    "Transpose": "Transpose",
}

---

三、第二步：图优化

GE 会对计算图做优化，包括常量折叠、算子融合、死代码消除。

导出优化前后的图

import torch
import torch_npu

# 加载模型
model = torch.jit.load("model.pt", map_location="npu:0")

# 导出优化前的图
import os
os.environ["GE_GRAPH_SAVE_PATH"] = "./before_opt"
model(torch.randn(1, 3, 224, 224).npu())

# 导出优化后的图
os.environ["GE_GRAPH_SAVE_PATH"] = "./after_opt"

常量折叠示例

# 优化前：运行时计算
class ModelBefore(torch.nn.Module):
    def forward(self, x):
        scale = torch.tensor(2.0)
        bias = torch.tensor(1.0)
        return x * scale + bias

# 优化后：编译时计算
# GE 会把 scale 和 bias 融合成一个常量
# 运行时只需要一次乘加操作

查看优化后的节点数量

# 编译时输出优化信息
atc --model=model.onnx \
    --framework=5 \
    --output=model \
    --log=info 2>&1 | grep "graph optimize"

# 输出示例：
# [GE] Before optimization: 178 nodes
# [GE] After optimization: 125 nodes
# [GE] Fused 53 nodes

---

四、第三步：内存规划

ATC 会分析每个 tensor 的生命周期，规划显存分配。

内存规划代码示例

# 模拟 GE 的内存规划算法
class MemoryPlanner:
    def __init__(self, total_memory):
        self.total_memory = total_memory
        self.allocations = {}  # tensor_name -> (offset, size)
        self.free_list = [(0, total_memory)]
    
    def allocate(self, tensor_name, size, start_time, end_time):
        """分配内存给 tensor"""
        # 找到合适的空闲块
        for i, (free_start, free_size) in enumerate(self.free_list):
            if free_size >= size:
                # 分配
                offset = free_start
                self.allocations[tensor_name] = {
                    "offset": offset,
                    "size": size,
                    "start": start_time,
                    "end": end_time
                }
                
                # 更新空闲列表
                if free_size > size:
                    self.free_list[i] = (free_start + size, free_size - size)
                else:
                    self.free_list.pop(i)
                
                return offset
        
        raise RuntimeError("Out of memory")
    
    def free(self, tensor_name):
        """释放 tensor 的内存"""
        alloc = self.allocations[tensor_name]
        offset, size = alloc["offset"], alloc["size"]
        self.free_list.append((offset, size))
        del self.allocations[tensor_name]

查看内存规划结果

# 编译时输出内存规划日志
export GE_MEMORY_PLANNING_LOG=1

atc --model=model.onnx \
    --framework=5 \
    --output=model

日志输出示例：

[MemoryPlanner] Tensor: conv1.output, Size: 0.5MB, Offset: 0
[MemoryPlanner] Tensor: bn1.output, Size: 0.5MB, Offset: 0 (reused)
[MemoryPlanner] Tensor: relu1.output, Size: 0.5MB, Offset: 0 (reused)
[MemoryPlanner] Total memory: 2.1MB
[MemoryPlanner] Peak memory: 1.2MB

---

五、第四步：算子选型

同一个算子有多种实现，ATC 会选择最优的。

Cube vs Vector 选择

# ATC 内部的算子选择逻辑（简化）
def select_matmul_impl(M, N, K):
    """选择矩阵乘的实现"""
    # Cube 实现要求矩阵大小 >= 16
    if M >= 16 and N >= 16 and K >= 16:
        if K % 16 == 0:  # K 需要对齐
            return "cube"
    
    # 否则用 Vector 实现
    return "vector"

# 测试
print(select_matmul_impl(64, 64, 64))   # 输出: cube
print(select_matmul_impl(8, 64, 64))    # 输出: vector
print(select_matmul_impl(64, 64, 17))   # 输出: vector (K 不对齐)

强制选择实现

import torch
import torch_npu

# 强制使用 Cube 实现
torch.npu.set_op_impl("matmul", impl="cube")
output = torch.matmul(a, b)

# 强制使用 Vector 实现
torch.npu.set_op_impl("matmul", impl="vector")
output = torch.matmul(a, b)

---

六、第五步：代码生成

ATC 生成二进制指令，打包进 .om 文件。

.om 文件结构

# .om 文件的结构（简化）
class OmFile:
    def __init__(self):
        self.header = OmHeader()
        self.model_data = ModelData()
        self.weight_data = WeightData()
        self.memory_plan = MemoryPlan()
        self.kernel_binaries = []

class OmHeader:
    magic: bytes          # 文件魔数
    version: int          # 版本号
    model_offset: int     # 模型数据偏移
    weight_offset: int    # 权重数据偏移
    memory_offset: int    # 内存规划偏移

class ModelData:
    graph: ComputeGraph   # 计算图
    op_descs: list        # 算子描述

class MemoryPlan:
    tensor_offsets: dict  # tensor 名称 -> 显存偏移
    total_size: int       # 总显存需求

读取 .om 文件信息

import struct

def parse_om_header(file_path):
    """解析 .om 文件头部"""
    with open(file_path, "rb") as f:
        # 读取魔数
        magic = f.read(4)
        
        # 读取版本
        version = struct.unpack("I", f.read(4))[0]
        
        # 读取偏移量
        model_offset = struct.unpack("Q", f.read(8))[0]
        weight_offset = struct.unpack("Q", f.read(8))[0]
        
        print(f"Magic: {magic}")
        print(f"Version: {version}")
        print(f"Model Offset: {model_offset}")
        print(f"Weight Offset: {weight_offset}")

parse_om_header("resnet50.om")

---

七、动态 Batch 编译

如果 batch size 需要变化，可以指定动态范围：

# 动态 batch 编译
atc --model=model.onnx \
    --framework=5 \
    --output=model_dynamic \
    --dynamic_batch_size="1,2,4,8,16,32"

动态 Shape 编译

# 动态序列长度
atc --model=bert.onnx \
    --framework=5 \
    --output=bert_dynamic \
    --input_shape_range="input:[1,1~512]"

Python 中使用动态模型

import torch
import torch_npu

# 加载动态模型
model = torch.jit.load("model_dynamic.om", map_location="npu:0")

# 不同 batch size 都可以用
for batch in [1, 4, 8, 16]:
    input_tensor = torch.randn(batch, 3, 224, 224).npu()
    output = model(input_tensor)
    print(f"Batch {batch}: output shape = {output.shape}")

---

八、编译失败排查

查看不支持的算子

# 编译并输出详细日志
atc --model=model.onnx \
    --framework=5 \
    --output=model \
    --log=debug 2>&1 | grep -i "not support"

# 输出示例：
# [ERROR] Op "CustomOp" is not supported
# [ERROR] Node "custom_node_1" has unsupported op

查看算子支持列表

# 查询 CANN 支持的算子
import os

opp_path = os.environ.get("ASCEND_HOME", "/usr/local/Ascend")
opp_config = f"{opp_path}/opp/built-in/op_impl/ai_core/tbe/config"

# 列出所有支持的算子
for root, dirs, files in os.walk(opp_config):
    for file in files:
        if file.endswith(".ini"):
            print(f"Supported op config: {file}")

替换不支持的算子

import torch

# 原始模型使用 HardSigmoid
class OriginalModel(torch.nn.Module):
    def forward(self, x):
        return torch.nn.functional.hardsigmoid(x)

# 替换为标准算子组合
class ReplacementModel(torch.nn.Module):
    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)

# 导出 ONNX
model = ReplacementModel()
torch.onnx.export(model, torch.randn(1, 3, 224, 224), "model.onnx")

---

九、编译性能对比

编译时间统计

time atc --model=resnet50.onnx --framework=5 --output=resnet50

# 输出示例：
# real    0m25.3s
# user    0m18.2s
# sys     0m4.1s

不同配置的编译时间

模型	静态 Shape	动态 Batch (6档)	动态 Shape (范围)
ResNet50	25s	78s	120s
BERT-base	45s	150s	280s
YOLOv5	30s	95s	-

---

参考资源

• ATC 编译参数详解：https://www.hiascend.com/document/detail/zh/CANN/
• ONNX 模型导出：https://pytorch.org/docs/stable/onnx.html
• 算子支持列表：https://www.hiascend.com/document/detail/zh/CANN/
• ATC 故障排查指南：https://www.hiascend.com/document/detail/zh/CANN/

---

总结

ATC 编译分五步：图解析把 ONNX 转成 GE 内部格式、图优化做融合和消除、内存规划分配显存、算子选型选择最优实现、代码生成打包 .om。每一步都可以通过参数控制：--log 输出详细日志、--enable_fusion 开关融合、--dynamic_batch_size 编译动态模型。遇到不支持的算子，先查算子列表，再用标准算子替换。动态编译时间比静态长 3-5 倍，开发阶段建议用静态，部署时再根据需要编译动态版本。