功能说明

APEX混合精度模块是一个集优化性能、精度收敛于一身的综合优化库。APEX功能模式介绍如下：

表1 混合精度模块功能
功能	开启方式举例	描述
O1配置模式	model, optimizer = amp.initialize(model, optimizer, opt_level="O1")	白名单：使用Cube加速的算子，总是使用float16计算，例如Conv2d、Matmul。黑名单：对精度有要求的算子，总是使用float32计算，例如Softmax、BN。其他：按照当前输入类型运算，例如Relu，MaxPool。
O2配置模式	model, optimizer = amp.initialize(model, optimizer, opt_level="O2")	针对全网中float32数据类型的算子，按照内置优化策略，自动将部分float32的算子降低精度到float16，从而在精度损失很小的情况下提升性能并减少内存使用。
O3配置模式	model, optimizer = amp.initialize(model, optimizer, opt_level="O3")	全部算子使用float16计算。
静态Loss Scale功能	model, optimizer = amp.initialize(model, optimizer, opt_level="O2",loss_scale=128.0)	用户可在混合精度训练过程中使用自定义的固定Loss Scale系数。
动态Loss Scale功能	model, optimizer = amp.initialize(model, optimizer, opt_level="O2",loss_scale="dynamic")	用户可使能在混合精度训练过程中，根据浮点计算异常状态动态调整Loss Scale系数。

推荐优先使用opt_level='O2', loss_scale=128.0的配置进行amp.initialize。若无法收敛，推荐使用opt_level='O1', loss_scale=128.0的配置进行amp.initialize。若依然无法收敛，推荐使用opt_level='O1', loss_scale=None的配置进行amp.initialize。

约束与限制

APEX当前版本的实现方式为python实现，不支持APEX中的自定义优化CUDA Kernel。
APEX当前版本只支持适配昇腾AI处理器的混合精度计算和多种融合优化器功能，其他功能暂未支持。融合优化器与原生优化器算法保持一致，但运算速度更快。

安装方法

请用户参考《CANN 软件安装指南》中“安装PyTorch”章节内的“安装APEX混合精度模块”，安装PyTorch框架对应的APEX即可。

使用说明

单卡训练场景。

导入混合精度模块。
```
from apex import amp
```

在模型和优化器定义之后初始化APEX模块。

model = ...
optimizer = ...
model, optimizer = amp.initialize(model, optimizer, combine_grad=True)

改写梯度反向传播loss.backward()。

loss = criterion(…) 
#将loss.backward()替换为如下形式
with amp.scale_loss(loss, optimizer) as scaled_loss:     
    scaled_loss.backward() 
optimizer.step()

多卡训练场景。

（推荐）使用APEX混合精度并关闭combine_ddp开关，并开启PyTorch框架自带的DistributedDataParallel（DDP）模式，DDP模式可参考PyTorch官方文档使用。

from torch.nn.parallel import DistributedDataParallel as DDP
from apex import amp
...
model = ...
optimizer = ...
#在模型、优化器定义之后，初始化APEX模块。
model, optimizer = amp.initialize(model, optimizer, combine_ddp=False)

使用APEX混合精度并开启combine_ddp开关，并关闭PyTorch框架自带的DDP模式。

from apex import amp
...
model = ...
optimizer = ...
#在模型、优化器定义之后，初始化APEX模块。
model, optimizer = amp.initialize(model, optimizer, combine_ddp=True)

更多混合精度模块的使用可参见官方文档。

APEX特性支持

替换亲和优化器函数。

适配后的APEX针对adadelta/adam/sgd/lamb做了昇腾AI处理器亲和性优化，得到的NPU融合优化器与原生算法保持一致，但运算速度更快。使用时只需将原有优化器替换为apex.optimizers.***，***为优化器名称，例如NpuFusedSGD。

样例原代码：

optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)    
model, optimizer = amp.initialize(model, optimizer, opt_level='O2', loss_scale=32.0)

修改后代码：

optimizer = apex.optimizers.NpuFusedSGD(model.parameters(), lr=args.lr, momentum=args.momentum)    
model, optimizer = amp.initialize(model, optimizer, opt_level='O2', loss_scale=32.0, combine_grad=True)

APEX使用小技巧：若希望通过自定义设置部分算子为float16提升性能或设置部分算子为float32提升精度，可以通过如下方法实现。

amp.register_half_function(torch, 'bmm') #填入算子API名称。以bmm为例，bmm会强制使用half进行计算 
amp.register_float_function(torch, 'bmm') #填入算子API名称。以bmm为例，bmm会强制使用float进行计算

当optimizer的可更新参数分为不同的多组，每组使用不同的策略，这种场景下，融合优化器只能优化掉第一个model分组的参数。因此，当优化器的参数相同时，尽量合并分组。

优化前：

optimizer = apex.optimizers.NpuFusedSGD([
             {'params': model.sharedNet.parameters()},
             {'params': model.bottleneck.parameters()},
             {'params': model.domain_classifier.parameters()},
             {'params': model.dcis.parameters()},
             {'params': model.source_fc.parameters(), 'lr': LEARNING_RATE},
         ], lr=LEARNING_RATE / 10, momentum=args.momentum, weight_decay=args.l2_decay)

优化后：

optimizer = apex.optimizers.NpuFusedSGD([
             {'params': list(model.sharedNet.parameters()) + list(model.bottleneck.parameters()) +  list(model.domain_classifier.parameters())+list(model.dcis.parameters())},
             {'params': model.source_fc.parameters(), 'lr': LEARNING_RATE},
         ], lr=LEARNING_RATE / 10, momentum=args.momentum, weight_decay=args.l2_decay)

开关分布式训练性能。
适配后的APEX针对数据并行场景做了昇腾AI处理器亲和性优化，支持利用融合grad进行加速，同时保持计算逻辑一致性。通过开启combine_ddp开关，也就是将amp.initialize()接口参数combine_ddp设置为True并关闭框架的DDP模式，即可开启该功能。
```
model, optimizer = amp.initialize(model, optimizer, combine_ddp=True) #配置运算加速参数combine_ddp
```
- 当前版本的实现方式主要为python实现，不支持AscendCL或者CUDA优化。
- 当前昇腾AI设备暂不支持原始APEX的FusedLayerNorm接口模块，如果模型原始脚本文件使用了FusedLayerNorm接口模块，需要在模型迁移过程中将脚本头文件from apex.normalization import FusedLayerNorm替换为from torch.nn import LayerNorm。
- 如果在DDP初始化后更改优化器，会导致主卡生效，副卡不生效，从而使模型出错。为了规避这个问题，需要将优化器更改放到DDP初始化前。