昇腾社区首页
中文
注册

混合精度

概述

基于NPU芯片的架构特性,模型运算会涉及到混合精度,即混合使用float16和float32数据类型的应用场景。使用float16代替float32有如下好处:

  • 对于中间变量的内存占用更少,节省内存的使用。
  • 因内存使用会减少,所以数据传出的时间也会相应减少。
  • float16的计算单元可以提供更快的计算性能。

但是,混合精度训练受限于float16表达的精度范围,单纯将float32转换成float16会影响训练收敛情况。为了保证部分计算使用float16来进行加速的同时能保证训练收敛,这里采用混合精度模块APEX来达到以上效果。混合精度模块APEX是一个集优化性能、精度收敛于一身的综合优化库。

安装模块

混合精度模块的安装请参考CANN 软件安装指南安装开发环境>在昇腾设备上安装>安装深度学习框架>安装PyTorch章节。

特性支持

混合精度模块功能和优化描述如表1所示。

表1 混合精度模块功能

功能

描述

O1配置模式

Conv、Matmul等使用float16计算,其他如Softmax、BN使用float32计算。

O2配置模式

针对全网中float32数据类型的算子,按照内置优化策略,自动将部分float32的算子降低精度到float16,从而在精度损失很小的情况下提升性能并减少内存使用。

静态Loss Scale功能

静态设置参数确保混合精度训练收敛。

动态Loss Scale功能

动态计算Loss Scale值并判断是否溢出。

当前版本的实现方式主要为Python实现,不支持AscendCL或者CUDA优化。

混合精度模型初始化

  • 使用APEX混合精度模块
    1. 从APEX库中导入AMP。
      from apex import amp
    2. 初始化AMP,使其能对模型、优化器以及PyTorch内部函数进行必要的改动。
      model, optimizer = amp.initialize(model, optimizer, opt_level="O2")

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

  • 使用框架自带AMP功能(PyTorch 1.8.1版本及以上)
    model = ...
    optimizer = ...
    #创建缩放器
    for epoch in epochs:
        for input, target in data:
            optimizer.zero_grad()
            with autocast():
                output = model(input)
                loss = loss_fn(output, target)
          ......

混合精度推理

按混合精度模型初始化后,正常执行模型正向计算即可。