简介

AOE（Ascend Optimization Engine）是一款自动调优工具，通过生成调优策略、编译和在运行环境上验证的闭环反馈机制，不断迭代出更优的调优策略，最终得到最佳的调优策略，从而可以更充分利用硬件资源，不断提升网络的性能，达到最优的效果。

如果要训练的模型为动态shape场景，不建议使用AOE调优。

调优流程

以PyTorch官网提供的Imagenet数据集训练脚本main.py为例，说明在PyTorch 1.8及以上版本通过AOE进行调优的方法。

以下步骤默认已配置CANN软件基础环境变量及Python环境变量，请在调优前提前配置。

对训练脚本进行迁移并训练，保证进行AOE调优前模型能执行成功并且精度达到预期。

首先需要配置自动迁移的环境变量。

export PYTHONPATH={CANN包安装目录}/ascend-toolkit/latest/tools/ms_fmk_transplt/torch_npu_bridge:$PYTHONPATH

在该训练脚本中导入以下Python包。

import torch_npu
import transfer_to_npu
from torch_npu.npu import amp

添加混合精度模块。

原训练脚本中进行loss计算及梯度更新的代码位于train()函数中，代码如下：

    # switch to train mode
    model.train()
    end = time.time()
    for i, (images, target) in enumerate(train_loader):
        # measure data loading time
        data_time.update(time.time() - end)
        # move data to the same device as model
        images = images.to(device, non_blocking=True)
        target = target.to(device, non_blocking=True)
        # compute output
        output = model(images)
        loss = criterion(output, target)
        # measure accuracy and record loss
        acc1, acc5 = accuracy(output, target, topk=(1, 5))
        losses.update(loss.item(), images.size(0))
        top1.update(acc1[0], images.size(0))
        top5.update(acc5[0], images.size(0))
        # compute gradient and do SGD step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()
        if i % args.print_freq == 0:
            progress.display(i + 1)

添加混合精度模块后，代码如下：

    # switch to train mode
    model.train()
    end = time.time()    
    scaler = amp.GradScaler()    
    for i, (images, target) in enumerate(train_loader):
        # measure data loading time
        data_time.update(time.time() - end)
        # move data to the same device as model
        images = images.to(device, non_blocking=True)
        target = target.to(device, non_blocking=True)
        # compute output
        with amp.autocast():
            output = model(images)
            loss = criterion(output, target)
        # measure accuracy and record loss
        acc1, acc5 = accuracy(output, target, topk=(1, 5))
        losses.update(loss.item(), images.size(0))
        top1.update(acc1[0], images.size(0))
        top5.update(acc5[0], images.size(0))
        # compute gradient and do SGD step
        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()
        if i % args.print_freq == 0:
            progress.display(i + 1)

设置好相关训练环境变量后，执行如下命令在0卡上进行训练，其中[imagenet-folder with train and val folders]需用户自行配置数据集路径。
```
python3 main.py -a resnet50 [imagenet-folder with train and val folders] --gpu 0
```

dump算子图。

在添加混合精度模块后的代码中添加如下代码使能算子dump。

    # switch to train mode
    model.train()
    end = time.time()    
    scaler = amp.GradScaler()
    # 这里的dump_path需用户自行指定，本例以训练脚本main.py所在路径下的dump文件夹为例
    # 例如可以设置为torch_npu.npu.set_aoe('./dump/')
    torch_npu.npu.set_aoe(dump_path)
    for i, (images, target) in enumerate(train_loader):
        # 仅需要运行一个step即可
        if i > 0:                      
            exit()
        # measure data loading time
        data_time.update(time.time() - end)
        # move data to the same device as model
        images = images.to(device, non_blocking=True)
        target = target.to(device, non_blocking=True)
        # compute output
        with amp.autocast():
            output = model(images)
            loss = criterion(output, target)
        # measure accuracy and record loss
        acc1, acc5 = accuracy(output, target, topk=(1, 5))
        losses.update(loss.item(), images.size(0))
        top1.update(acc1[0], images.size(0))
        top5.update(acc5[0], images.size(0))
        # compute gradient and do SGD step
        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()
        if i % args.print_freq == 0:
            progress.display(i + 1)

拉起训练命令执行训练。会在dump_path路径下生成算子图文件。

（可选）指定调优后自定义知识库的存储路径。
```
export TUNE_BANK_PATH=/home/HwHiAiUser/custom_tune_bank
```
如不配置此环境变量，自定义知识库默认存储在${HOME}/Ascend/latest/data/aoe/custom/op/${soc_version}路径下。自定义知识库可以应用于推理和训练的性能提升，使用方法请参考《AOE工具使用指南》中的“附录>如何使用调优后的自定义知识库”章节。

执行调优（调优时间较长，请耐心等待）。

# dump_path为上述步骤中配置的算子图文件存放路径
# 以训练脚本main.py所在路径下的dump文件夹为例
aoe --job_type=2 --model_path=dump_path

调优完成后，会在/home/HwHiAiUser/custom_tune_bank/${soc_version}路径下生成自定义知识库。

调优命令参数说明可参考下表。

表1 命令参数说明
参数名称	参数简述
---model_path	原始模型文件路径，该路径下可以存放多个模型文件。此处配置为dump算子图步骤中的dump_path。
-j --job_type	调优模式。取值为： 1：子图调优，当前PyTorch框架不支持。 2：算子调优

调优完成后，设置如下环境变量并将使能算子dump的代码删除。
```
export ENABLE_TUNE_BANK=True
```
重新执行训练。

AOE自动调优

简介

调优流程