Pytorch mixed precision 概述（混合精度）

點擊上方“機器學習與生成對抗網絡”，關注星標

獲取有趣、好玩的前沿干貨！

作者：知乎 superjie13

https://www.zhihu.com/people/superjie13

本文對pytorch中的mixed precision進行測試。主要包括兩部分，第一部分為mixed precision使用概述，第二部分為實際測試。參考torch官網 Automatic Mixed Precision

01

Mixed precision使用概述

通常，automatic mixed precision training 需要使用 torch.cuda.amp.autocast 和 torch.cuda.amp.GradScaler 。

1. 1 首先實例化 torch.cuda.amp.autocast(enable=True) 作為上下文管理器或者裝飾器，從而使腳本使用混合精度運行。注意：autocast 一般情況下只封裝前向傳播過程（包括loss的計算），并不包括反向傳播（反向傳播的數據類型與相應前向傳播中的數據類型相同）。

1. 2 使用Gradient scaling 防止在反向傳播過程由于中梯度太?。╢loat16無法表示小幅值的變化）從而下溢為0的情況。torch.cuda.amp.GradScaler() 可以自動進行gradient scaling。注意：由于GradScaler()對gradient進行了scale，因此每個參數的gradient應該在optimizer更新參數前unscaled，從而使學習率不受影響。

import torchvisionimport torchimport torch.cuda.ampimport gcimport time
# Timing utilitiesstart_time = None
def start_timer():    global start_time    gc.collect()    torch.cuda.empty_cache()    torch.cuda.reset_max_memory_allocated()    torch.cuda.synchronize()  # 同步后得出的時間才是實際運行的時間    start_time = time.time()
def end_timer_and_print(local_msg):    torch.cuda.synchronize()    end_time = time.time()    print("\n" + local_msg)    print("Total execution time = {:.3f} sec".format(end_time - start_time))    print("Max memory used by tensors = {} bytes".format(torch.cuda.max_memory_allocated()))
num_batches = 50batch_size = 70epochs = 3
# 隨機創(chuàng)建訓練數據data = [torch.randn(batch_size, 3, 224, 224, device="cuda") for _ in range(num_batches)]targets = [torch.randint(0, 1000, size=(batch_size, ), device='cuda') for _ in range(num_batches)]# 創(chuàng)建一個模型net = torchvision.models.resnext50_32x4d().cuda()# 定義損失函數loss_fn = torch.nn.CrossEntropyLoss().cuda()# 定義優(yōu)化器opt = torch.optim.SGD(net.parameters(), lr=0.001)
# 是否使用混合精度訓練use_amp = True
# Constructs scaler once, at the beginning of the convergence run, using default args.# If your network fails to converge with default GradScaler args, please file an issue.# The same GradScaler instance should be used for the entire convergence run.# If you perform multiple convergence runs in the same script, each run should use# a dedicated fresh GradScaler instance.  GradScaler instances are lightweight.scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
start_timer()for epoch in range(epochs):    for input, target in zip(data, targets):        with torch.cuda.amp.autocast(enabled=use_amp):            output = net(input)            loss = loss_fn(output, target)        # 放大loss  Calls backward() on scaled loss to create scaled gradients.        scaler.scale(loss).backward()
        # scaler.step() first unscales the gradients of the optimizer's assigned params.        # If these gradients do not contain infs or NaNs, optimizer.step() is then called,        # otherwise, optimizer.step() is skipped.        scaler.step(opt)
        # Updates the scale for next iteration.        scaler.update()        opt.zero_grad(set_to_none=True) # set_to_none=True here can modestly improve performanceend_timer_and_print("Mixed precision:")

02
混合精度測試

測試環(huán)境：ubuntu18.04, pytorch 1.7.1, python3.7, RTX2080-8G

2.1 use_amp = False

batch size = 40

2.2 use_amp = True

batch size = 40

從實驗2.1和2.2中，可以發(fā)現在batch size=40的情況下，不使用混合精度時，GPU內存占用為7011MB，運行時間為47.55 s。而使用混合精度時，GPU內存占用為4997MB，運行時間為27.006 s。在當前運行配置中，內存占用節(jié)省了約28.73%，運行時間節(jié)省了約43.21%。這也就意味著我們可以使用更大的batch size來提升運行效率。

2.3 use_amp = True

batch size = 70

猜您喜歡：

等你著陸！【GAN生成對抗網絡】知識星球！

超100篇！CVPR 2020最全GAN論文梳理匯總！

附下載 | 《Python進階》中文版

附下載 | 經典《Think Python》中文版

附下載 | 《Pytorch模型訓練實用教程》

附下載 | 最新2020李沐《動手學深度學習》

附下載 | 《可解釋的機器學習》中文版

附下載 |《TensorFlow 2.0 深度學習算法實戰(zhàn)》

附下載 | 超100篇！CVPR 2020最全GAN論文梳理匯總！

附下載 |《計算機視覺中的數學方法》分享