更快的計算，更高的內存效率：PyTorch混合精度模型AMP介紹

時間 2020-03-22

標籤更快計算高的內存效率 pytorch 混合精度模型 amp 介紹简体版

原文原文鏈接

做者：Rahul Agarwal python

您是否知道反向傳播算法是Geoffrey Hinton在1986年的《天然》雜誌上提出的？ git

一樣的，卷積網絡由Yann le cun於1998年首次提出，並進行了數字分類，他使用了單個卷積層。直到2012年下半年，Alexnet才經過使用多個卷積層在imagenet上實現最早進的技術來推廣卷積網絡。 github

那麼，是什麼讓他們直到如今才這麼出名？算法

只有在咱們擁有大量計算資源的狀況下，咱們才能進行足夠的實驗並充分利用深度學習的潛力。網絡

可是咱們真的充分利用了現有的計算資源嗎？咱們能夠作得更好嗎？性能

這篇文章是關於利用Tensor Cores和自動混合精度來更快地訓練深度學習網絡的。學習

什麼是Tensor Cores？

根據NVIDIA網站的介紹：測試

NVIDIA Turing和Volta GPU由Tensor Cores提供支持。Tensor Cores是一項革命性技術，可提供開創性的AI性能。 Tensor Core能夠加速AI核心的大型矩陣運算，並在單個運算中執行混合精度矩陣乘法和累加計算。在一個NVIDIA GPU中並行運行數百個Tensor Core，這能夠極大地提升吞吐量和運行效率。

簡單地說; 它們是**專用於特定類型矩陣操做的專用內核。** 網站

咱們能夠將兩個FP16矩陣相乘並將其添加到FP16 / FP32矩陣中，從而獲得FP16 / FP32矩陣。 Tensor Core支持混合精度數學，即輸入爲半精度（FP16），輸出爲全精度（FP32）。這樣的操做對於許多深度學習任務具備內在的價值，而且Tensor Core爲該操做提供了專用的硬件。 spa

與FP32相比FP16主要有一下兩個好處。

1.FP16須要較少的內存，所以更易於訓練和部署大型神經網絡。它還減小了數據移動。

2.使用Tensor Core，大大提升了數學運算的速度但下降了精度。 NVIDIA提供的Volta GPU的確切數量是：FP16中爲125 TFlops，而FP32中爲15.7 TFlops（加速8倍）

可是也有缺點。從FP32轉到FP16時，必然會下降精度。

FP32與FP16：FP32具備八個指數位和23個小數位，而FP16具備五個指數位和十個小數位。

可是真的須要FP32嗎？

FP16實際上能夠很好地表示大多數權重和漸變色。所以，存儲和使用FP32所需的這些額外數位只是浪費。

那麼，咱們應該怎麼使用Tensor Cores呢？

我檢查了個人Titan RTX GPU。它擁有576個Tensor Cores以及4,608個NVIDIA CUDA內核。可是如何使用這些Tensor Cores？

坦白地說，NVIDIA能夠輕鬆地將Tensor Cores與自動混合精度一塊兒使用，並提供了幾行代碼。咱們須要在代碼中作兩件事：

1.將FP32所需的操做（如Softmax）分配給FP32，而將FP16能夠完成的操做（如卷積）自動分配給FP16。

2.使用損失縮放來保留較小的梯度值。梯度值可能超出FP16的範圍。在這種狀況下，將對梯度值進行縮放，使其保持在FP16範圍內。

就算您還不瞭解背景細節也能夠。由於它的代碼實現相對簡單。

使用PyTorch進行混合精度訓練

讓咱們從PyTorch中的基本網絡開始。

N, D_in, D_out = 64, 1024, 512
x = torch.randn(N, D_in, device="cuda")
y = torch.randn(N, D_out, device="cuda")
model = torch.nn.Linear(D_in, D_out).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)for to in range(500):
y_pred = model(x)
loss = torch.nn.functional.mse_loss(y_pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()

想要利用自動混合精度訓練的優點，咱們首先須要安裝apex庫。只需在終端中運行如下命令。

$ git clone https://github.com/NVIDIA/apex
$ cd apex
$ pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

而後，咱們只需在神經網絡代碼中添加幾行便可利用自動混合精度（AMP）。

from apex import amp
N, D_in, D_out = 64, 1024, 512
x = torch.randn(N, D_in, device="cuda")
y = torch.randn(N, D_out, device="cuda")
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 to in range(500):
y_pred = model(x)
loss = torch.nn.functional.mse_loss(y_pred, y)
optimizer.zero_grad()
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
optimizer.step()

在這裏，您能夠看到咱們使用amp初始化了模型。咱們還經過使用amp.scale_loss指定損失縮放比例。

### 標杆分析法

咱們可使用這個很棒的存儲庫對放大器的性能進行測試，該存儲庫對CIFAR數據集上的VGG16模型進行測試。我只須要更改幾行代碼便可爲咱們工做。您能夠在此處找到修改後的版本。要本身運行測試代碼，您可能須要：

git clone https://github.com/MLWhiz/data_science_blogs
cd data_science_blogs/amp/pytorch-apex-experiment/
python run_benchmark.py
python make_plot.py --GPU 'RTX' --method 'FP32''FP16''amp' --batch 12825651210242048

這將在主目錄中爲您生成如下圖形：