昇腾AOE调优引擎
AOE 调优引擎
AOE 概述
AOE(Ascend Optimization Engine)是 CANN 提供的自动化性能调优引擎,通过搜索最优的算子执行参数(如 Tiling 策略)来提升模型性能,无需开发者手动调优。
AOE 支持三种调优模式:
| 模式 | 全称 | 调优对象 | 典型收益 |
|---|---|---|---|
| OPAT | Operator Auto Tuning | 单算子 Tiling 参数 | 10-30% |
| SGAT | Subgraph Auto Tuning | 子图级融合与调度 | 15-40% |
| GDAT | Gradient Auto Tuning | 梯度计算优化 | 5-20% |
AOE 工作原理
┌─────────────────────────────────────────────────────────┐
│ AOE 调优流程 │
│ │
│ 1. 收集算子信息 │
│ ├── 算子类型、输入输出形状、数据类型 │
│ └── 当前执行性能基线 │
│ │
│ 2. 搜索空间构建 │
│ ├── Tiling 参数空间(分块大小、循环顺序等) │
│ └── 融合策略空间 │
│ │
│ 3. 自动搜索(贝叶斯优化 / 遗传算法) │
│ ├── 在真实硬件上执行候选配置 │
│ └── 记录性能数据 │
│ │
│ 4. 生成调优知识库 │
│ └── 保存最优配置到 .json 文件 │
│ │
│ 5. 推理/训练时加载调优结果 │
│ └── 自动使用最优配置执行 │
└─────────────────────────────────────────────────────────┘
OPAT:算子自动调优
通过 ATC 触发 OPAT
# 在模型转换时同时进行算子调优
atc --model=model.onnx \--framework=5 \--output=model_tuned \--soc_version=Ascend910B3 \--input_shape="input:1,3,224,224" \--enable_scope_fusion_passes=true \--op_compiler_cache_mode=enable \--op_compiler_cache_dir=./op_cache \--aoe_mode=1 # 1=OPAT
独立运行 AOE 工具
# 方式一:基于模型文件调优
aoe --framework=5 \--model=model.onnx \--job_type=1 \--output=./aoe_result \--soc_version=Ascend910B3 \--input_shape="input:1,3,224,224"# 方式二:基于已有 .om 文件调优
aoe --om=model.om \--job_type=1 \--output=./aoe_result \--soc_version=Ascend910B3
在 Python 中触发 OPAT
import torch
import torch_npu# 开启 AOE 调优模式
torch.npu.set_option({"ACL_OP_COMPILER_CACHE_MODE": "enable","ACL_OP_COMPILER_CACHE_DIR": "./op_cache"
})# 运行几个 step 收集算子信息
for i, batch in enumerate(dataloader):if i >= 10: # 收集 10 个 step 的数据breakoutput = model(batch)# 触发 AOE 调优
torch_npu.npu.set_aoe("./aoe_result")
SGAT:子图自动调优
SGAT 在子图级别进行优化,寻找最优的算子融合策略。
# 子图调优
aoe --framework=5 \--model=model.onnx \--job_type=2 \--output=./sgat_result \--soc_version=Ascend910B3 \--input_shape="input:1,3,224,224" \--tuning_time_limit=3600 # 调优时间限制(秒)
子图调优配置
// sgat_config.json
{"job_type": "subgraph","tuning_time_limit": 3600,"max_tuning_count": 100,"search_algorithm": "bayesian","target_metric": "throughput"
}
GDAT:梯度自动调优
GDAT 专门针对训练场景的梯度计算进行优化。
# 梯度调优(需要训练数据)
aoe --framework=5 \--model=model.onnx \--job_type=3 \--output=./gdat_result \--soc_version=Ascend910B3 \--input_shape="input:1,3,224,224;label:1" \--loss_name="loss"
加载调优结果
ATC 转换时加载
atc --model=model.onnx \--framework=5 \--output=model_optimized \--soc_version=Ascend910B3 \--input_shape="input:1,3,224,224" \--op_compiler_cache_mode=enable \--op_compiler_cache_dir=./aoe_result # 指向调优结果目录
运行时加载
import torch_npu# 加载 AOE 调优结果
torch.npu.set_option({"ACL_OP_COMPILER_CACHE_MODE": "enable","ACL_OP_COMPILER_CACHE_DIR": "./aoe_result"
})
调优结果分析
调优报告
AOE 会生成详细的调优报告:
aoe_result/
├── summary.json # 调优摘要
├── op_tuning/
│ ├── MatMul_fp16.json # MatMul 算子调优结果
│ ├── Conv2D_fp16.json # Conv2D 算子调优结果
│ └── ...
└── profiling/├── before_tuning/ # 调优前性能数据└── after_tuning/ # 调优后性能数据
summary.json 示例:
{"total_ops": 156,"tuned_ops": 89,"performance_improvement": "23.5%","tuning_time": "1823s","top_improved_ops": [{"op": "MatMul_1", "improvement": "45.2%"},{"op": "Conv2D_3", "improvement": "31.8%"},{"op": "LayerNorm_2", "improvement": "28.4%"}]
}
调优最佳实践
1. 选择代表性输入
# 使用实际业务中最常见的输入形状进行调优
# 避免使用极端形状(太小或太大)
--input_shape="input:8,3,224,224" # 使用实际 batch size
2. 分阶段调优
# 第一阶段:快速 OPAT(1-2小时)
aoe --job_type=1 --tuning_time_limit=3600# 第二阶段:深度 SGAT(4-8小时)
aoe --job_type=2 --tuning_time_limit=28800
3. 调优结果复用
# 相同模型结构、相同硬件的调优结果可以复用
# 将调优结果纳入版本管理
git add aoe_result/
git commit -m "Add AOE tuning results for ResNet50 on 910B3"
4. 增量调优
# 模型更新后,只对变化的算子重新调优
aoe --model=model_v2.onnx \--job_type=1 \--output=./aoe_result_v2 \--base_result=./aoe_result # 基于已有结果增量调优
性能调优效果示例
| 模型 | 调优前(ms) | 调优后(ms) | 提升 |
|---|---|---|---|
| ResNet50 | 3.2 | 2.4 | 25% |
| BERT-Base | 12.5 | 9.8 | 22% |
| YOLOv5s | 8.1 | 6.2 | 24% |
| GPT-2 | 45.3 | 35.1 | 22% |
以上数据为示意,实际效果因硬件和模型而异
