原標題：GPU捉襟見肘還想訓練大批量模型？誰說不可以

選自Medium

深度學習模型和數據集的規模增長速度已經讓 GPU 算力也開始捉襟見肘，如果你的 GPU 連一個樣本都容不下，你要如何訓練大批量模型？通過本文介紹的方法，我們可以在訓練批量甚至單個訓練樣本大于 GPU 內存時，在單個或多個 GPU 服務器上訓練模型。

2018 年的大部分時間我都在試圖訓練神經網絡時克服 GPU 極限。無論是在含有 1.5 億個參數的語言模型(如 OpenAI 的 或最近類似的 模型)還是饋入 3000 萬個元素輸入的元學習神經網絡(如我們在一篇 ICLR 論文《》中提到的模型)，我都只能在 GPU 上處理很少的訓練樣本。

但在多數情況下，隨機梯度下降算法需要很大批量才能得出不錯的結果。

如果你的 GPU 只能處理很少的樣本，你要如何訓練大批量模型？

有幾個工具、技巧可以幫助你解決上述問題。在本文中，我將自己用過、學過的東西整理出來供大家參考。

在這篇文章中，我將主要討論 PyTorch 框架。有部分工具尚未包括在 PyTorch(1.0 版本)中，因此我也寫了自定義代碼。

我們將著重探討以下問題：

在訓練批量甚至單個訓練樣本大于 GPU 內存，要如何在單個或多個 GPU 服務器上訓練模型；

如何盡可能高效地利用多 GPU 機器；

在分布式設備上使用多個機器的最簡單訓練方法。

在一個或多個 GPU 上訓練大批量模型

你建的模型不錯，在這個簡潔的任務中可能成為新的 SOTA，但每次嘗試在一個批量處理更多樣本時，你都會得到一個 CUDA RuntimeError：內存不足。

這位網友指出了你的問題！

但你很確定將批量加倍可以優化結果。

你要怎么做呢？

這個問題有一個簡單的解決方法：梯度累積。

梯度下降優化算法的五個步驟。

與之對等的 PyTorch 代碼也可以寫成以下五行：

predictions = model(inputs) # Forward pass

loss = loss_function(predictions, labels) # Compute loss function

loss.backward() # Backward pass

optimizer.step() # Optimizer step

predictions = model(inputs) # Forward pass with new parameters

在 loss.backward() 運算期間，為每個參數計算梯度，并將其存儲在與每個參數相關聯的張量——parameter.grad 中。

累積梯度意味著，在調用 optimizer.step() 實施一步梯度下降之前，我們會對 parameter.grad 張量中的幾個反向運算的梯度求和。在 PyTorch 中這一點很容易實現，因為梯度張量在不調用 model.zero_grad() 或 optimizer.zero_grad() 的情況下不會重置。如果損失在訓練樣本上要取平均，我們還需要除以累積步驟的數量。

以下是使用梯度累積訓練模型的要點。在這個例子中，我們可以用一個大于 GPU 最大容量的 accumulation_steps 批量進行訓練：

model.zero_grad() # Reset gradients tensors

fori, (inputs, labels) inenumerate(training_set):

predictions = model(inputs) # Forward pass

loss = loss_function(predictions, labels) # Compute loss function

loss = loss / accumulation_steps # Normalize our loss (if averaged)

loss.backward() # Backward pass

if(i+ 1) % accumulation_steps == 0: # Wait for several backward steps

optimizer.step() # Now we can do an optimizer step

model.zero_grad() # Reset gradients tensors

if(i+ 1) % evaluation_steps == 0: # Evaluate the model when we...

evaluate_model() # ...have no gradients accumulated

擴展到極致

你可以在 GPU 上訓練連一個樣本都無法加載的模型嗎？

如果你的架構沒有太多跳過連接，這就是可能的！解決方案是使用梯度檢查點(gradient-checkpointing)來節省計算資源。

基本思路是沿著模型將梯度在小組件中進行反向傳播，以額外的前饋傳遞為代價，節約存儲完整的反向傳播圖的內存。這個方法比較慢，因為我們需要添加額外的計算來減少內存要求，但在某些設置中挺有意思，比如在非常長的序列上訓練 RNN 模型(示例參見 https://medium.com/huggingface/from-zero-to-research-an-introduction-to-meta-learning-8e16e677f78a)。

這里不再贅述，讀者可以查看以下鏈接：

TensorFlow：https://github.com/openai/gradient-checkpointing

PyTorch 文檔：https://pytorch.org/docs/stable/checkpoint.html

「節約內存」(Memory-poor)策略需要 O(1) 的內存(但是要求 O(n2) 的計算步)。

充分利用多 GPU 機器

現在我們具體來看如何在多 GPU 上訓練模型。

在多 GPU 服務器上訓練 PyTorch 模型的首選策略是使用 torch.nn.DataParallel。該容器可以在多個指定設備上分割輸入，按照批維度(batch dimension)分割，從而實現模塊應用的并行化。

DataParallel 非常容易使用，我們只需添加一行來封裝模型：

parallel_model = torch.nn.DataParallel(model) # Encapsulate the model

predictions = parallel_model(inputs) # Forward pass on multi-GPUs

loss = loss_function(predictions, labels) # Compute loss function

loss.backward() # Backward pass

optimizer.step() # Optimizer step

predictions = parallel_model(inputs) # Forward pass with new parameters

但是，DataParallel 有一個問題：GPU 使用不均衡。

在一些設置下，GPU-1 會比其他 GPU 使用率高得多。

這個問題從何而來呢？下圖很好地解釋了 DataParallel 的行為：

使用 torch.nn.DataParallel 的前向和后向傳播。

在前向傳播的第四步(右上)，所有并行計算的結果都聚集在 GPU-1 上。這對很多分類問題來說是件好事，但如果你在大批量上訓練語言模型時，這就會成為問題。

我們可以快速計算語言模型輸出的大小：

語言模型輸出中的元素數量。

假設我們的數據集有 4 萬詞匯，每一條序列有 250 個 token、每個 batch 中有 32 條序列，那么序列中的每一個元素需要 4 個字節的內存空間，模型的輸出大概為 1.2GB。要儲存相關的梯度張量，我們就需要把這個內存翻倍，因此我們的模型輸出需要 2.4GB 的內存。

這是典型 10GB GPU 內存的主要部分，意味著相對于其它 GPU，GPU - 1 會被過度使用，從而限制了并行化的效果。

如果不調整模型和／或優化方案，我們就無法輕易減少輸出中的元素數量。但我們可以確保內存負載在 GPU 中更均勻地分布。

多 GPU 機器上的均衡負載

解決辦法是把每部分輸出保留在其 GPU 上，而不是將它們聚集到 GPU-1 上。我們也需要分配損失標準計算，計算損失并進行反向傳播。

幸而，張航開源了一個名為 PyTorch-Encoding 的 PyTorch 包，它包含了這些定制的并行化功能。

我提取并稍稍改動了這個模塊，你可以從以下地址下載 gist(parallel.py)來納入并調用你的代碼。它主要包括兩個模塊：DataParallelModel 和 DataParallelCriterion，它們的用途如下：

下載地址：https://gist.github.com/thomwolf/7e2407fbd5945f07821adae3d9fd1312

fromparallel importDataParallelModel, DataParallelCriterion

parallel_model = DataParallelModel(model) # Encapsulate the model

parallel_loss = DataParallelCriterion(loss_function) # Encapsulate the loss function

predictions = parallel_model(inputs) # Parallel forward pass

# "predictions" is a tuple of n_gpu tensors

loss = parallel_loss(predictions, labels) # Compute loss function in parallel

loss.backward() # Backward pass

optimizer.step() # Optimizer step

predictions = parallel_model(inputs) # Parallel forward pass with new parameters

DataParallelModel 和 torch.nn.DataParallel 的區別在于，前向傳播的輸出(predictions)沒有聚集在 GPU-1 上，而是作為 n_gpu 張量的元組，每個張量分布在相應的 GPU 上。

DataParallelCriterion 容器封裝了損失函數，并把 n_gpu 張量元組和目標標簽張量作為輸入。它在每個 GPU 上并行計算損失函數，像 DataParallel 分割模型輸入一樣分割目標標簽張量。

下圖說明了 DataParallelModel/DataParallelCriterion 的內部情況：

使用 DataParallelModel 和 DataParallelCriterion。

以下是你可能會遇到的兩個特定案例的解決辦法：

你的模型輸出幾個張量：你可能想分解它們：output_1, output_2 = zip(*predictions)

有時候你并不想使用并行損失函數：收集 CPU 上的所有張量：gathered_predictions = parallel.gather(predictions)

分布式訓練：在多臺機器上訓練

在更大的批量上訓練時，我們要如何控制多個服務器的算力呢？

最簡單的選擇是使用 PyTorch 的 DistributedDataParallel，它幾乎可以說是以上討論的 DataParallel 的直接替代元件。

但要注意：盡管代碼看起來很相似，但在分布式設定中訓練模型要改變工作流程，因為你必須在每個節點上啟動一個獨立的 Python 訓練腳本。正如我們將看到的，一旦啟動，這些訓練腳本可以通過使用 PyTorch 分布式后端一起同步化。

在實踐中，這意味著每個訓練腳本將擁有：

它自己的優化器，并在每次迭代中執行一個完整的優化步驟，不需要進行參數傳播(DataParallel 中的步驟 2)；

一個獨立的 Python 解釋器：這也將避免 GIL-freeze，這是在單個 Python 解釋器上驅動多個并行執行線程時會出現的問題。

當多個并行前向調用由單個解釋器驅動時，在前向傳播中大量使用 Python 循環/調用的模型可能會被 Python 解釋器的 GIL 放慢速度。通過這種設置，DistributedDataParallel 甚至在單臺機器設置中也能很方便地替代 DataParallel。

現在我們直接討論代碼和用途。

DistributedDataParallel 是建立在 torch.distributed 包之上的，這個包可以為同步分布式運算提供低級原語，并能以不同的性能使用多種后端(tcp、gloo、mpi、nccl)。在這篇文章中，我將選擇一種簡單的開箱即用的方式來使用它，但你應該閱讀文檔和 Séb Arnold 寫的教程來深入理解這個模塊。

文檔：https://pytorch.org/docs/stable/distributed.html

教程：https://pytorch.org/tutorials/intermediate/dist_tuto.html

我們將考慮使用具有兩個 4 - GPU 服務器(節點)的簡單但通用的設置：

主服務器(服務器 1)擁有一個可訪問的 IP 地址和一個用于通信的開放端口。

改寫 Python 訓練腳本以適應分布式訓練

首先我們需要改寫腳本，從而令其可以在每臺機器(節點)上獨立運行。我們將實現完全的分布式訓練，并在每個節點的每塊 GPU 上運行一個獨立的進程，因此總共需要 8 個進程。

我們的訓練腳本有點長，因為需要為同步化初始化分布式后端，封裝模型并準備數據，以在數據的一個子集上來訓練每個進程(每個進程都是獨立的，因此我們需要自行處理)。以下是更新后的代碼：

fromtorch.utils.data.distributed importDistributedSampler

fromtorch.utils.data importDataLoader

# Each process runs on 1 GPU device specified by the local_rank argument.

parser = argparse.ArgumentParser()

parser.add_argument( "--local_rank", type=int)

args = parser.parse_args()

# Initializes the distributed backend which will take care of sychronizing nodes/GPUs

torch.distributed.init_process_group(backend= 'nccl')

# Encapsulate the model on the GPU assigned to the current process

device = torch.device( 'cuda', arg.local_rank)

model = model.to(device)

distrib_model = torch.nn.parallel.DistributedDataParallel(model,

device_ids=[args.local_rank],

output_device=args.local_rank)

# Restricts data loading to a subset of the dataset exclusive to the current process

sampler = DistributedSampler(dataset)

dataloader = DataLoader(dataset, sampler=sampler)

forinputs, labels indataloader:

predictions = distrib_model(inputs.to(device)) # Forward pass

loss = loss_function(predictions, labels.to(device)) # Compute loss function

loss.backward() # Backward pass

optimizer.step() # Optimizer step

啟動 Python 訓練腳本的多個實例

我們就快完成了，只需要在每個服務器上啟動訓練腳本的一個實例。

為了運行腳本，我們將使用 PyTorch 的 torch.distributed.launch 工具。它將用來設置環境變量，并用正確的 local_rank 參數調用每個腳本。

第一臺機器是最主要的，它應該對于所有其它機器都是可訪問的，因此擁有一個可訪問的 IP 地址(我們的案例中是 192.168.1.1)以及一個開放端口(在我們的案例中是 1234)。在第一臺機器上，我們使用 torch.distributed.launch 來運行訓練腳本：

python -m torch.distributed.launch --nproc_per_node=4 --nnodes=2 --node_rank=0 --master_addr= "192.168.1.1"--master_port=1234 OUR_TRAINING_.py (--arg1 --arg2 --arg3 and all other arguments of our training ) # Optimizer step

在第二臺機器上，我們類似地啟動腳本：

python -m torch.distributed.launch --nproc_per_node=4 --nnodes=2 --node_rank=1 --master_addr= "192.168.1.1"--master_port=1234 OUR_TRAINING_.py (--arg1 --arg2 --arg3 and all other arguments of our training )

這兩個命令是相同的，除了—node_rank 參數，其在第一臺機器上被設為 0，在第二臺機器上被設為 1(如果再加一臺機器，則設為 2，以此類推…)。

點擊閱讀原文報名 INTERFACE，探索微軟小冰全雙工語音對話技術。

責任編輯：

總結

以上是生活随笔為你收集整理的python 调用gpu算力_GPU捉襟见肘还想训练大批量模型？谁说不可以的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。