快捷導(dǎo)航

AMP?Tensor?Cores節(jié)省內(nèi)存PyTorch模型詳解

更新時間：2022年10月24日 11:14:58 作者：ronghuaiyang

這篇文章主要為大家介紹了AMP?Tensor?Cores節(jié)省內(nèi)存PyTorch模型詳解，有需要的朋友可以借鑒參考下，希望能夠有所幫助，祝大家多多進(jìn)步，早日升職加薪

導(dǎo)讀

只需要添加幾行代碼，就可以得到更快速，更省顯存的PyTorch模型。

你知道嗎，在1986年Geoffrey Hinton就在Nature論文中給出了反向傳播算法？

此外，卷積網(wǎng)絡(luò)最早是由Yann le cun在1998年提出的，用于數(shù)字分類，他使用了一個卷積層。但是直到2012年晚些時候，Alexnet才通過使用多個卷積層來實(shí)現(xiàn)最先進(jìn)的imagenet。

那么，是什么讓他們現(xiàn)在如此出名，而不是之前呢？

只有在我們擁有大量計算資源的情況下，我們才能夠在最近的過去試驗(yàn)和充分利用深度學(xué)習(xí)的潛力。

但是，我們是否已經(jīng)足夠好地使用了我們的計算資源呢？我們能做得更好嗎？

這篇文章的主要內(nèi)容是關(guān)于如何利用Tensor Cores和自動混合精度更快地訓(xùn)練深度學(xué)習(xí)網(wǎng)絡(luò)。

什么是Tensor Cores?

根據(jù)NVIDIA的網(wǎng)站：

NVIDIA Turing和Volta GPUs都是由Tensor Cores驅(qū)動的，這是一項突破性的技術(shù)，提供了突破性的AI性能。Tensor Cores可以加速AI核心的大矩陣運(yùn)算，在一次運(yùn)算中就可以完成混合精度的矩陣乘法和累加運(yùn)算。在一個NVIDIA GPU上有數(shù)百個Tensor Cores并行運(yùn)行，這大大提高了吞吐量和效率。

簡單地說，它們是專門的cores，非常適合特定類型的矩陣操作。

我們可以將兩個FP16矩陣相乘，并將其添加到一個FP16/FP32矩陣中，從而得到一個FP16/FP32矩陣。Tensor cores支持混合精度數(shù)學(xué)，即以半精度(FP16)進(jìn)行輸入，以全精度(FP32)進(jìn)行輸出。上述類型的操作對許多深度學(xué)習(xí)任務(wù)具有內(nèi)在價值，而Tensor cores為這種操作提供了專門的硬件。

現(xiàn)在，使用FP16和FP32主要有兩個好處。

FP16需要更少的內(nèi)存，因此更容易訓(xùn)練和部署大型神經(jīng)網(wǎng)絡(luò)。它還只需要較少的數(shù)據(jù)移動。
數(shù)學(xué)運(yùn)算在降低精度的Tensor cores運(yùn)行得更快。NVIDIA給出的Volta GPU的確切數(shù)字是：FP16的125 TFlops vs FP32的15.7 TFlops(8倍加速)。

但也有缺點(diǎn)。當(dāng)我們從FP32轉(zhuǎn)到FP16時，我們需要降低精度。

FP32 vs FP16: FP32 有8個指數(shù)位和23個分?jǐn)?shù)位，而FP16有5個指數(shù)位和10個分?jǐn)?shù)位。

但是FP32真的有必要嗎？

實(shí)際上，F(xiàn)P16可以很好地表示大多數(shù)權(quán)重和梯度。所以存儲和使用FP32是很浪費(fèi)的。

那么，我們?nèi)绾问褂肨ensor Cores？

我檢查了一下我的Titan RTX GPU有576個tensor cores和4608個NVIDIA CUDA核心。但是我如何使用這些tensor cores呢？

坦白地說，NVIDIA用幾行代碼就能提供自動混合精度，因此使用tensor cores很簡單。我們需要在代碼中做兩件事：

需要用到FP32的運(yùn)算比如Softmax之類的就分配用FP32，而Conv之類的操作可以用FP16的則被自動分配用FP16。

使用損失縮放 為了保留小的梯度值。梯度值可能落在FP16的范圍之外。在這種情況下，梯度值被縮放，使它們落在FP16范圍內(nèi)。

如果你還不了解背景細(xì)節(jié)也沒關(guān)系，代碼實(shí)現(xiàn)相對簡單。

使用PyTorch進(jìn)行混合精度訓(xùn)練：

讓我們從PyTorch中的一個基本網(wǎng)絡(luò)開始。

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()

為了充分利用自動混合精度訓(xùn)練的優(yōu)勢，我們首先需要安裝apex庫。只需在終端中運(yùn)行以下命令。

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

然后，我們只需向神經(jīng)網(wǎng)絡(luò)代碼中添加幾行代碼，就可以利用自動混合精度(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.initialize初始化了我們的模型。我們還使用amp.scale_loss來指定損失縮放。

基準(zhǔn)測試

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 128 256 512 1024 2048

這會在home目錄中生成下面的圖: