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在行程內部生效，

標籤：其他

上一篇：Mask R-CNN

下一篇：用Python構建個性化智能鬧鐘