pytorch 訓練_在PyTorch中使用混合精度訓練

簡介

PyTorch在進行模型訓練時采用的都是單精度(FP32)訓練的方法,，也就是每個資料都占用4個位元組(32bit)的空間。半精度(FP16)資料則隻占用2個位元組(16bit)的存儲空間。

是以FP16相比FP32節省了一半的存儲空間和位寬，不僅在相同顯存占用的情況下容納更多樣本，而且可以加快計算速度。而在多數情況下，FP16帶來的精度降低對模型性能的損失很小，并且有許多可行的辦法可以進一步降低這種損失。

半精度訓練對于RTX系列顯示卡的提升更加明顯。因為RTX系列顯示卡配備了針對FP16運算專門優化的Tensor Core，在半精度訓練下速度提升非常明顯。

由于PyTorch 1.4及以前的版本本身不支援半精度訓練，是以需要安裝NVIDIA針對PyTorch的APEX拓展。

針對半精度訓練帶來的精度下降和溢出問題，APEX主要采用了以下兩種方法解決問題。

1. 混合精度訓練(Mixed Precision) 混合精度訓練的精髓在于“在記憶體中用FP16做儲存和乘法進而加速計算，用FP32做累加避免舍入誤差”。混合精度訓練的政策有效地緩解了舍入誤差的問題。
2. 損失放大(Loss Scaling) 即使用了混合精度訓練，還是會存在無法收斂的情況，原因是激活梯度的值太小，造成了下溢出(Underflow)。損失放大的思路是：
   1. 反向傳播前，将損失變化(dLoss)手動增大 

      

       倍，是以反向傳播時得到的中間變量(激活函數梯度)則不會溢出；
   2. 反向傳播後，将權重梯度縮 

      

       倍，恢複正常值。

從PyTorch 1.5開始，将自身支援自動混合精度訓練。但是目前官網下載下傳的PyTorch 1.5穩定版中的混合精度訓練子產品是不完整的，需要手動将你的資料轉成FP16。是以最好下載下傳并使用PyTorch預覽版(PyTorch 1.6)。

APEX安裝使用(PyTorch穩定版)

安裝

git clone https://github.com/NVIDIA/apexcd apexpython setup.py install
           

若下載下傳速度太慢，可以将倉庫位址更換為https://gitee.com/kabu1204/apex

使用

使用非常簡單，隻需在原先代碼的基礎上做三處改動

1.從apex中導入amp子產品

from apex import amp
           

2.初始化模型

model=nn().cuda()optimizer=optim.Adam(...)model, optimizer = amp.initialize(net, optimizer, opt_level="O1")  #O1中的O是大寫字母O
           

3.通過amp反向傳播

...loss = loss_fn(preds, labels)with amp.scale_loss(loss, optimizer) as scaled_loss:     scaled_loss.backward()#loss.backward()...
           

PyTorch預覽版的自動混合精度訓練

安裝PyTorch預覽版

打開PyTorch官方安裝頁面

建議使用Pip方式安裝，使用conda安裝速度很慢。

最好建立一個新的Anaconda環境，或者用原來的環境先解除安裝PyTorch。

驗證一下安裝成功

使用

使用同樣很簡單。官方給出的執行個體已經講得很清楚了。

# Creates model and optimizer in default precisionmodel = Net().cuda()optimizer = optim.SGD(model.parameters(), ...)# Creates a GradScaler once at the beginning of training.scaler = torch.cuda.amp.GradScaler()for epoch in epochs:    for input, target in data:        optimizer.zero_grad()        # Runs the forward pass with autocasting.        with torch.cuda.amp.autocast():            output = model(input)            loss = loss_fn(output, target)        # Scales loss.  Calls backward() on scaled loss to create scaled gradients.        # Backward passes under autocast are not recommended.        # Backward ops run in the same precision that autocast used for corresponding forward ops.        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(optimizer)        # Updates the scale for next iteration.        scaler.update()
           

速度對比

單精度	APEX	torch.cuda.amp
~1350 秒/EPOCH	~940 秒/EPOCH	~760 秒/EPOCH

用mini-ImageNet在darknet上測試

APEX比單精度快了約43.6%(所用時間縮短了約30%)

torch.cuda.amp比單精度快了約77.6%(所用時間縮短了約43%)

torch.cuda.amp比APEX快了約23.6%(所用時間縮短了約19%)

相關資料

APEX官方文檔

NVIDIA官方混合精度訓練文檔

PyTorch AUTOMATIC MIXED PRECISION EXAMPLES