AOE 自动调优：让模型跑得更快

50084

517人浏览 · 2026-05-24 12:10:33

50084 · 2026-05-24 12:10:33 发布

前言

同样的模型，不同输入 shape、不同硬件配置，最优的算子参数不一样。AOE（Auto Optimization Engine）做的事情就是搜索这些最优参数——比如矩阵乘的 tile 大小、卷积的循环展开方式。

一、AOE 能调什么

AOE 的调优对象是算子的实现参数，不是模型结构。具体包括：

参数类型	示例	影响
Tiling 参数	矩阵乘的 block 大小	决定并行度和 UB 利用率
循环展开	卷积的循环展开次数	影响 Cube/Vector 选择
内存布局	NC1HWC0 vs NCHW	影响带宽利用率
算子选择	Cube vs Vector 实现	影响延迟和吞吐

AOE 不会改变模型的数学逻辑，只是找到更高效的执行方式。

二、调优模式选择

AOE 支持两种调优模式：

1. GA（遗传算法）模式

适合调优空间大的场景，搜索时间较长（30-60 分钟）。

# 使用 GA 模式调优
aoe --model=model.onnx \
    --framework=5 \
    --job_type=2 \
    --mode=ga \
    --output=optimized_model

2. RL（强化学习）模式

适合快速调优，搜索时间较短（5-15 分钟），适合开发阶段快速迭代。

# 使用 RL 模式调优
aoe --model=model.onnx \
    --framework=5 \
    --job_type=2 \
    --mode=rl \
    --output=optimized_model

模式对比

模式	搜索时间	效果	适用场景
GA	30-60 分钟	最优	最终部署
RL	5-15 分钟	次优	开发迭代

三、ResNet50 调优实战

准备工作

# 1. 编译原始模型
atc --model=resnet50.onnx \
    --framework=5 \
    --output=resnet50_base

# 2. 准备调优数据
python prepare_data.py --output calibration_data.bin

启动调优

aoe --model=resnet50.onnx \
    --framework=5 \
    --job_type=2 \
    --mode=ga \
    --calibration_data=calibration_data.bin \
    --output=resnet50_optimized

调优过程输出

[AOE] Start tuning...
[AOE] Generation 1/50, best latency: 12.5ms
[AOE] Generation 10/50, best latency: 10.2ms
[AOE] Generation 20/50, best latency: 9.1ms
[AOE] Generation 30/50, best latency: 8.5ms
[AOE] Generation 40/50, best latency: 8.3ms
[AOE] Generation 50/50, best latency: 8.1ms
[AOE] Tuning completed!
[AOE] Original latency: 12.5ms
[AOE] Optimized latency: 8.1ms
[AOE] Improvement: 35.2%

调优结果

指标	调优前	调优后	提升
推理延迟	12.5ms	8.1ms	35%
AICore 利用率	62%	78%	+16%
带宽利用率	55%	71%	+16%

四、查看调优参数

调优完成后，AOE 会输出调优参数文件，可以用于后续编译。

# 调优参数保存在 .json 文件里
ls -la resnet50_optimized/
# output:
# resnet50_optimized.om
# tuning_result.json

调优参数示例

{
  "op_tuning": {
    "Conv_0": {
      "tiling_strategy": "cube_opt",
      "block_m": 16,
      "block_n": 16,
      "block_k": 32
    },
    "MatMul_10": {
      "impl_type": "cube",
      "tiling": {
        "M_tile": 64,
        "N_tile": 64,
        "K_tile": 128
      }
    }
  }
}

每个算子都有对应的调优参数，包括 tiling 策略、block 大小、实现类型等。

五、应用调优结果

调优后的 .om 文件可以直接使用。如果只想保存调优参数供后续编译使用：

# 使用调优参数编译新模型
atc --model=resnet50.onnx \
    --framework=5 \
    --output=resnet50_final \
    --load_tuning_result=resnet50_optimized/tuning_result.json

Python 中应用调优结果

import torch
import torch_npu

# 加载调优后的模型
model = torch.jit.load("resnet50_optimized.om", map_location="npu:0")

# 正常推理
output = model(input_tensor)

六、不同模型的调优效果

Transformer 模型（BERT-base）

aoe --model=bert_base.onnx \
    --framework=5 \
    --job_type=2 \
    --mode=ga \
    --output=bert_optimized

指标	调优前	调优后	提升
首token延迟	95ms	78ms	18%
生成速度	52 tok/s	61 tok/s	17%
显存占用	12GB	11.2GB	-7%

Transformer 的主要优化点是 FlashAttention 的 tile 大小和 Softmax 的实现选择。

目标检测模型（YOLOv5）

aoe --model=yolov5s.onnx \
    --framework=5 \
    --job_type=2 \
    --mode=rl \
    --output=yolov5_optimized

指标	调优前	调优后	提升
推理延迟	18ms	13ms	28%
后处理时间	4.5ms	1.2ms	73%

YOLOv5 的主要优化点是 NMS 算子的实现和卷积的 tiling 策略。

七、自定义调优空间

如果只想调优特定的算子，可以指定调优空间：

# 定义调优空间配置
tuning_config = {
    "ops_to_tune": ["Conv_0", "MatMul_10"],  # 只调这两个算子
    "tiling_range": {
        "block_m": [8, 16, 32, 64],
        "block_n": [8, 16, 32, 64],
        "block_k": [16, 32, 64, 128]
    },
    "impl_options": ["cube", "vector"]
}

import json
with open("tuning_config.json", "w") as f:
    json.dump(tuning_config, f)

# 使用自定义调优空间
aoe --model=model.onnx \
    --framework=5 \
    --job_type=2 \
    --tuning_config=tuning_config.json \
    --output=model_optimized

八、调优最佳实践

1. 选择正确的输入 shape

调优时的输入 shape 应该和生产环境一致：

# 生产环境输入是 (1, 3, 640, 640)
atc --model=model.onnx \
    --input_shape="input:1,3,640,640" \
    --framework=5 \
    --output=model

aoe --model=model.onnx \
    --framework=5 \
    --input_shape="input:1,3,640,640" \
    --job_type=2 \
    --output=model_optimized

如果输入 shape 变化，调优参数可能不再最优，需要重新调优。

2. 批量调优多个 shape

# 调优多个 batch size
for batch in 1 4 8 16; do
    aoe --model=model.onnx \
        --framework=5 \
        --input_shape="input:${batch},3,224,224" \
        --job_type=2 \
        --output=model_batch${batch}
done

3. 增量调优

如果只改了模型的一部分，可以只调优改动部分：

# 加载之前的调优结果，只调优新增算子
aoe --model=model_v2.onnx \
    --framework=5 \
    --job_type=2 \
    --load_tuning_result=model_v1/tuning_result.json \
    --incremental \
    --output=model_v2_optimized

参考资源

AOE 调优指南：https://www.hiascend.com/document/detail/zh/CANN/
调优参数详解：https://www.hiascend.com/document/detail/zh/CANN/
AOE 样例仓库：https://atomgit.com/cann/samples
性能调优最佳实践：https://www.hiascend.com/document/detail/zh/CANN/

总结

AOE 的调优本质是搜索算子的最优参数——tile 大小、实现方式、内存布局。GA 模式搜索彻底，适合最终部署；RL 模式搜索快速，适合开发迭代。调优效果因模型而异：ResNet50 能提 30%，Transformer 能提 15-20%。调优时的输入 shape 要和生产一致，shape 变化需要重新调优。调优结果可以保存成 .json 文件，后续编译直接加载，不需要重新搜索。