解決Pytorch半精度浮點型網絡訓練的問題
用Pytorch1.0進行半精度浮點型網絡訓練需要註意下問題:
1、網絡要在GPU上跑,模型和輸入樣本數據都要cuda().half()
2、模型參數轉換為half型,不必索引到每層,直接model.cuda().half()即可
3、對於半精度模型,優化算法,Adam我在使用過程中,在某些參數的梯度為0的時候,更新權重後,梯度為零的權重變成瞭NAN,這非常奇怪,但是Adam算法對於全精度數據類型卻沒有這個問題。
另外,SGD算法對於半精度和全精度計算均沒有問題。
還有一個問題是不知道是不是網絡結構比較小的原因,使用半精度的訓練速度還沒有全精度快。這個值得後續進一步探索。
對於上面的這個問題,的確是網絡很小的情況下,在1080Ti上半精度浮點型沒有很明顯的優勢,但是當網絡變大之後,半精度浮點型要比全精度浮點型要快。
但具體快多少和模型的大小以及輸入樣本大小有關系,我測試的是要快1/6,同時,半精度浮點型在占用內存上比較有優勢,對於精度的影響尚未探究。
將網絡再變大些,epoch的次數也增大,半精度和全精度的時間差就表現出來瞭,在訓練的時候。
補充:pytorch半精度,混合精度,單精度訓練的區別amp.initialize
看代碼吧~
mixed_precision = True try: # Mixed precision training https://github.com/NVIDIA/apex from apex import amp except: mixed_precision = False # not installed model, optimizer = amp.initialize(model, optimizer, opt_level='O1', verbosity=1)
為瞭幫助提高Pytorch的訓練效率,英偉達提供瞭混合精度訓練工具Apex。號稱能夠在不降低性能的情況下,將模型訓練的速度提升2-4倍,訓練顯存消耗減少為之前的一半。
文檔地址是:https://nvidia.github.io/apex/index.html
該 工具 提供瞭三個功能,amp、parallel和normalization。由於目前該工具還是0.1版本,功能還是很基礎的,在最後一個normalization功能中隻提供瞭LayerNorm層的復現,實際上在後續的使用過程中會發現,出現問題最多的是pytorch的BN層。
第二個工具是pytorch的分佈式訓練的復現,在文檔中描述的是和pytorch中的實現等價,在代碼中可以選擇任意一個使用,實際使用過程中發現,在使用混合精度訓練時,使用Apex復現的parallel工具,能避免一些bug。
默認訓練方式是 單精度float32
import torch model = torch.nn.Linear(D_in, D_out) optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) for img, label in dataloader: out = model(img) loss = LOSS(out, label) loss.backward() optimizer.step() optimizer.zero_grad()
半精度 model(img.half())
import torch model = torch.nn.Linear(D_in, D_out).half() optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) for img, label in dataloader: out = model(img.half()) loss = LOSS(out, label) loss.backward() optimizer.step() optimizer.zero_grad()
接下來是混合精度的實現,這裡主要用到Apex的amp工具。
代碼修改為:
加上這一句封裝,
model, optimizer = amp.initialize(model, optimizer, opt_level=“O1”)
import torch model = torch.nn.Linear(D_in, D_out).cuda() optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) model, optimizer = amp.initialize(model, optimizer, opt_level="O1") for img, label in dataloader: out = model(img) loss = LOSS(out, label) # loss.backward() with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() optimizer.step() optimizer.zero_grad()
實際流程為:調用amp.initialize按照預定的opt_level對model和optimizer進行設置。在計算loss時使用amp.scale_loss進行回傳。
需要註意以下幾點:
在調用amp.initialize之前,模型需要放在GPU上,也就是需要調用cuda()或者to()。
在調用amp.initialize之前,模型不能調用任何分佈式設置函數。
此時輸入數據不需要在轉換為半精度。
在使用混合精度進行計算時,最關鍵的參數是opt_level。他一共含有四種設置值:‘00’,‘01’,‘02’,‘03’。實際上整個amp.initialize的輸入參數很多:
但是在實際使用過程中發現,設置opt_level即可,這也是文檔中例子的使用方法,甚至在不同的opt_level設置條件下,其他的參數會變成無效。(已知BUG:使用‘01’時設置keep_batchnorm_fp32的值會報錯)
概括起來:
00相當於原始的單精度訓練。01在大部分計算時采用半精度,但是所有的模型參數依然保持單精度,對於少數單精度較好的計算(如softmax)依然保持單精度。02相比於01,將模型參數也變為半精度。
03基本等於最開始實驗的全半精度的運算。值得一提的是,不論在優化過程中,模型是否采用半精度,保存下來的模型均為單精度模型,能夠保證模型在其他應用中的正常使用。這也是Apex的一大賣點。
在Pytorch中,BN層分為train和eval兩種操作。
實現時若為單精度網絡,會調用CUDNN進行計算加速。常規訓練過程中BN層會被設為train。Apex優化瞭這種情況,通過設置keep_batchnorm_fp32參數,能夠保證此時BN層使用CUDNN進行計算,達到最好的計算速度。
但是在一些fine tunning場景下,BN層會被設為eval(我的模型就是這種情況)。此時keep_batchnorm_fp32的設置並不起作用,訓練會產生數據類型不正確的bug。此時需要人為的將所有BN層設置為半精度,這樣將不能使用CUDNN加速。
一個設置的參考代碼如下:
def fix_bn(m): classname = m.__class__.__name__ if classname.find('BatchNorm') != -1: m.eval().half() model.apply(fix_bn)
實際測試下來,最後的模型準確度上感覺差別不大,可能有輕微下降;時間上變化不大,這可能會因不同的模型有差別;顯存開銷上確實有很大的降低。
以上為個人經驗,希望能給大傢一個參考,也希望大傢多多支持WalkonNet。
推薦閱讀:
- pytorch 使用半精度模型部署的操作
- PyTorch梯度裁剪避免訓練loss nan的操作
- Pytorch 中的optimizer使用說明
- pytorch 實現L2和L1正則化regularization的操作
- 使用Pytorch實現two-head(多輸出)模型的操作