自动混合精度(AMP)

混合精度(Mix Precision)训练是指在训练时,对网络不同的部分采用不同的数值精度,对追求速度的算子(比如 conv、matmul)可以用较低精度(比如 float16)从而获得明显的性能提升,而对其它更看重精度的算子(比如 log、softmax)则保留较高的精度(比如 float32)。

得益于 NVIDIA TensorCore 的存在(需要 Volta, Turing, Ampere 架构 GPU),对于 conv、matmul 占比较多的网络,混合精度训练一般能使网络整体的训练速度有较大的提升(2-3X)。

接口介绍

在 MegEngine 中,使用 autocast 接口可以实现对网络中相关 op 数据类型的自动转换:

import numpy as np
import megengine as mge
from megengine import amp
from megengine.hub import load
net = load("megengine/models", "resnet18", pretrained=False)
inp = mge.tensor(np.random.normal(size=(1, 3, 224, 224)), dtype="float32")

with amp.autocast():    # 使用 autocast context 接口
    oup = net(inp)
print(oup.dtype)

上面样例中是把 autocast 作为 context manager 使用,也可以使用装饰器的写法:

@amp.autocast()
def train_func(inp):    # 使用 autocast 装饰器
    oup = net(inp)
    return oup

oup = train_func(inp)
print(oup.dtype)

或者使用单独的开关:

amp.enabled = True
oup = net(inp)
amp.enabled = False
print(oup.dtype)

在开启 autocast 后,网络中间结果的数值类型会变成 float16,其对应的梯度自然也是 float16。而由于 float16 的数值范围比 float32 要小,所以如果遇到特别小的数(比如 loss、gradient),float16 就难以精确表达,这时候一般需要进行梯度放缩(Gradient Scaling)。做法是对网络的 loss 进行放大,从而使反向传播时网络中间结果对应的梯度也得到相同的放大,减少精度的损失,而梯度反传到参数时,仍会是 float32 的类型,在等比缩小之后,并不会影响参数的更新。

在 MegEngine 中使用梯度放缩,可以通过 GradScaler 接口进行。

import megengine.functional as F
from megengine.autodiff import GradManager
from megengine.optimizer import SGD

gm = GradManager().attach(net.parameters())
opt = SGD(net.parameters(), lr=0.01)
scaler = amp.GradScaler()           # 使用 GradScaler 进行梯度缩放

image = mge.tensor(np.random.normal(size=(1, 3, 224, 224)), dtype="float32")
label = mge.tensor(np.zeros(1), dtype="int32")

@amp.autocast()
def train_step(image, label):
    with gm:
        logits = net(image)
        loss = F.nn.cross_entropy(logits, label)
        scaler.backward(gm, loss)   # 通过 GradScaler 修改反传行为
    opt.step().clear_grad()
    return loss

train_step(image, label)

上面的样例中,通过替换 gm.backward(loss)scaler.backward(gm, loss),可以实现对 loss、gradient 的自动放缩,实际上包含三个步骤:

所以如果需要更加精细的操作,比如累积多个 iter 的梯度,那么可以使用以下等价形式:

@amp.autocast()
def train_step(image, label):
    with gm:
        logits = net(image)
        loss = F.nn.cross_entropy(logits, label)
        gm.backward(loss, dy=mge.tensor(scaler.scale_factor))   # 对应步骤一
    # 这里可以插入对梯度的自定义操作
    scaler.unscale(gm.attached_tensors())                       # 对应步骤二
    scaler.update()                                             # 对应步骤三
    opt.step().clear_grad()
    return loss

train_step(image, label)

我们可以形象地把上面两种方式分别称为自动挡和手动挡。

通过以上接口,就可以在无需修改模型代码的条件下,只修改训练代码实现混合精度训练,大幅提升网络的训练速度了。