NPU亲和优化器替换调优案例

问题背景

在PyTorch混合精度模式下，每次迭代执行一次参数更新时，在Loss乘/除以缩放系数的过程中都会包含连续多个小算子（如add、mul、sqrt等）下发，由于小算子在NPU上计算快，导致算子在CPU上的下发成为性能的主要瓶颈。本调优案例以MobileNetV1模型为例，进行梯度融合调优解决此性能问题。

总体思路

使用昇腾AI处理器亲和的融合优化器替换原始优化器，本例中使用torch_npu.optim.NpuFusedSGD替换原始优化器torch.optim.SGD。昇腾AI处理器亲和的融合优化器可以在梯度更新过程中将连续多个小算子融合后下发，提升模型计算性能，完整融合优化器替换表请参见NPU亲和优化器替换。

环境准备

点击获取链接，获取MobileNetV1模型，并根据readme准备训练环境和数据集。

从昇腾ModelZoo上获取的MobileNetV1模型脚本为已经进行了梯度融合优化的版本。若用户想使用该模型脚本体验本案例中的调优流程，可先将模型脚本main.py中的优化器定义代码改为原始优化器torch.optim.SGD，再执行以下流程。

瓶颈识别

1. 用户可在代码中搜索关键词optim以搜索当前模型的优化器函数名，对照表1查看是否有可替换的亲和优化器。
2. 修改训练脚本还原原生优化器函数，此操作在原生模型脚本中无需进行。

       model = model.to(device)
   
       # define loss function (criterion) and optimizer
       criterion = nn.CrossEntropyLoss().to(device)
   
       optimizer = torch.optim.SGD(model.parameters(), args.lr,
                                                   momentum=args.momentum,
                                               weight_decay=args.weight_decay)
   
       if args.apex:
           model, optimizer = amp.initialize(model, optimizer,
                                             opt_level=args.apex_opt_level,
                                             loss_scale=args.loss_scale_value,
                                             combine_grad=True)

3. 执行训练，记录耗时。

   bash ./test/train_full_1p.sh --data_path=/data/path/    # 请根据实际情况更改data_path

   耗时数据如所示。

   图1 原生优化器数据
   

性能调优

以下调优步骤基于已完成模型向NPU的迁移。

1. 在模型脚本开头添加库代码。

   import torch_npu
   import torch_npu.optim

2. 找到模型脚本main.py中的优化器定义代码，将原始优化器替换为对应的NPU亲和融合优化器。样例代码如下。
   原代码：

   optimizer = torch.optim.SGD(model.parameters(), args.lr,
                                           momentum=args.momentum,
                                           weight_decay=args.weight_decay)

   修改后代码：

   optimizer = torch_npu.optim.NpuFusedSGD(model.parameters(), args.lr,
                                           momentum=args.momentum,
                                           weight_decay=args.weight_decay)

   或

   optimizer = apex.optimizers.NpuFusedSGD(model.parameters(), args.lr,
                                           momentum=args.momentum,
                                           weight_decay=args.weight_decay)

   

   当前昇腾适配的MobileNetV1模型使用的是apex组件包中的NpuFusedSGD，此处修改为torch_npu.optim.NpuFusedSGD可能会导致精度异常。

3. 拉起模型训练，进行profiling，生成json文件。以下命令以单卡训练为例。

   bash ./test/train_full_1p.sh --data_path=/data/path/    # 请根据实际情况更改data_path

4. 执行训练，对比未设置亲和优化器时的训练数据，每秒处理的图片数提升，性能有一定程度优化。

   耗时数据如图2所示。

   图2 亲和优化器数据