Pytorch原生AMP支持使用方法(1.6版本)

AMP：Automatic mixed precision，自动混合精度，能够在神经网络推理过程当中，针对不一样的层，采用不一样的数据精度进行计算，从而实现节省显存和加快速度的目的。

在Pytorch 1.5版本及之前，经过NVIDIA出品的插件apex，能够实现amp功能。

从Pytorch 1.6版本之后，Pytorch将amp的功能吸取入官方库，位于torch.cuda.amp模块下。

本文为针对官方文档主要内容的简要翻译和本身的理解。

1. Introduction

torch.cuda.amp提供了对混合精度的支持。为实现自动混合精度训练，须要结合使用以下两个模块：

• torch.cuda.amp.autocast：autocast主要用做上下文管理器或者装饰器，来肯定使用混合精度的范围。
• torch.cuda.amp.GradScalar：GradScalar主要用来完成梯度缩放。

2. Typical Mixed Precision Training

一个典型的amp应用示例以下：

# 定义模型和优化器
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# 在训练最开始定义GradScalar的实例
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        # 利用with语句，在autocast实例的上下文范围内，进行模型的前向推理和loss计算
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        # 对loss进行缩放，针对缩放后的loss进行反向传播
        # （此部分计算在autocast()做用范围之外）
        scaler.scale(loss).backward()

        # 将梯度值缩放回原尺度后，优化器进行一步优化
        scaler.step(optimizer)

        # 更新scalar的缩放信息
        scaler.update()

3. Working with Unscaled Gradients

待更新

4. Working with Scaled Gradients

待更新

5. Working with Multiple Models, Losses, and Optimizers

若是模型的Loss计算部分输出多个loss，须要对每个loss值执行scaler.scale。

若是网络具备多个优化器，对任一个优化器执行scaler.unscale_，并对每个优化器执行scaler.step。

而scaler.update只在最后执行一次。

应用示例以下：

scaler = torch.cuda.amp.GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer0.zero_grad()
        optimizer1.zero_grad()
        with autocast():
            output0 = model0(input)
            output1 = model1(input)
            loss0 = loss_fn(2 * output0 + 3 * output1, target)
            loss1 = loss_fn(3 * output0 - 5 * output1, target)

        scaler.scale(loss0).backward(retain_graph=True)
        scaler.scale(loss1).backward()

        # 选择其中一个优化器执行显式的unscaling
        scaler.unscale_(optimizer0)
		# 对每个优化器执行scaler.step
        scaler.step(optimizer0)
        scaler.step(optimizer1)
		# 完成全部梯度更新后，执行一次scaler.update
        scaler.update()

6. Working with Multiple GPUs

针对多卡训练的状况，只影响autocast的使用方法，GradScaler的用法与以前一致。

6.1 DataParallel in a single process

在每个不一样的cuda设备上，torch.nn.DataParallel在不一样的进程中执行前向推理，而autocast只在当前进程中生效，所以，以下方式的调用是不生效的：

model = MyModel()
dp_model = nn.DataParallel(model)

# 在主进程中设置autocast
with autocast():
    # dp_model的内部进程并不会对autocast生效
    output = dp_model(input)
    # loss的计算在主进程中执行，autocast能够生效，但因为前面执行推理时已经失效，所以总体上是不正确的
    loss = loss_fn(output)

有效的调用方式以下所示：

# 方法1：在模型构建中，定义forwar函数时，采用装饰器方式
MyModel(nn.Module):
    ...
    @autocast()
    def forward(self, input):
       ...

# 方法2：在模型构建中，定义forwar函数时，采用上下文管理器方式
MyModel(nn.Module):
    ...
    def forward(self, input):
        with autocast():
            ...

# DataParallel的使用方式不变
model = MyModel().cuda()
dp_model = nn.DataParallel(model)

# 在模型执行推理时，因为前面模型定义时的修改，在各cuda设备上的子进程中autocast生效
# 在执行loss计算是，在主进程中，autocast生效
with autocast():
    output = dp_model(input)
    loss = loss_fn(output)

6.2 DistributedDataParallel, one GPU per process

torch.nn.parallel.DistributedDataParallel在官方文档中推荐每一个GPU执行一个实例的方法，以达到最好的性能表现。

在这种模式下，DistributedDataParallel内部并不会再启动子进程，所以对于autocast和GradScaler的使用都没有影响，与典型示例保持一致。

6.3 DistributedDataParallel, multiple GPUs per process

与DataParallel 的使用相同，在模型构建时，对forward函数的定义方式进行修改，保证autocast在进程内部生效。