增大Batch訓練神經網路：單GPU、多GPU及分散式配置的實用技巧

摘要： 2018年中的大部分時間，我都在嘗試利用訓練神經網路克服GPUs的侷限。無論是在包含1.5億個引數的語言模型中，比如 OpenAI’s huge Generative Pre-trained Transformer (or the recent and similar BERT model ...

2018年中的大部分時間，我都在嘗試利用訓練神經網路克服GPUs的侷限。無論是在包含1.5億個引數的語言模型中，比如 ofollow,noindex" target="_blank">OpenAI’s huge Generative Pre-trained Transformer (or the recent and similar BERT model )，還是在擁有3000萬個輸入元素的神經網路中，我都只能利用GPU處理很少的訓練樣本。

可是若想利用隨機梯度下降演算法得出不錯的結果，大批量的訓練樣本必不可少。

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

接下來，我將介紹幾類工具和技巧。

本文主要會討論PyTorch框架，並就以下幾個問題進行探討：

1. 當訓練批量甚至單個訓練樣本大於GPU記憶體時，如何訓練模型；
2. 如何高效地利用多GPU機器；
3. 如何在分散式裝置上簡單的使用多個機器。

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

你構建了一個不錯的模型，可在嘗試處理更多樣本時，卻得到CUDA RuntimeError：記憶體不足。

根據網友的回答你明白，加倍批量可以對結果進行優化。

此時， 梯度累積（accumulating gradients） 可以幫助到你。

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中。

梯度累積意味著，在呼叫potimizer.step()實現梯度下降之前，我們會求取parameter.grad張量中的幾個反向操作的梯度和。

如下是使用梯度累積訓練模型的示例。

model.zero_grad()# Reset gradients tensors
for i, (inputs, labels) in enumerate(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-checkpoingting） 節省計算資源。

梯度檢查點會將我們連續計算的元前饋和元反向傳播切分成片段。但由於需要增加額外的計算以減少記憶體需求，該方法效率不高。不過，它在某些示例中又有較為明顯的優勢，比如在長序列上訓練RNN模型， 點選此處 檢視詳情。

或有興趣可進入下列文件進行查詢：

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

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

A “Memory-poor” strategy that needs O(1) memory (but requires O(n²) computation steps) — From Yaroslav Bulatov’s nice post: https://medium.com/tensorflow/fitting-larger-networks-into-memory-583e3c758ff9

多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使用不均衡的問題，下圖給出了相應解釋：

Forward and Backward passes with torch.nn.DataParallel

在前向傳播的第四個步驟（見右上）中，GPU-1彙集了所有平行計算的結果。

通過下列所示的方式能夠計算出語言模型輸出的大小：

Number of elements in the output of a language model

現有如下假設：資料集共含4萬詞彙，序列中包含250 tokens，每個batch 包含32個示例，每個元素4 bytes，模型的輸出佔用1.2GB。但我們需要2.4GB的記憶體才能儲存相關的梯度張量。

這種儲存方式會使得GPU-1被過度使用，從而造成GPU使用不均衡的問題。

多GPU機器上的負載均衡

想要解決GPU使用不均衡的問題需要將每部分輸出都保留在原有的GPU上，而不彙集於GPU-1。

張航開源了名為PyTorch-Encoding的包，可用於緩解上述問題。

我對這個開源包做了一些調整，你可以點選此處下載 parallel.py 。此包中包含兩個模組：DataParallelModel以及DataParallelCriterion，如下所示：

from parallel import DataParallelModel, 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_gup張量的元組分佈在相應的GPU上。

DataParallelCriterion容器封裝了損失函式，並且將n_gpu張量的元組和目標標籤張量作為輸入。

下圖描述了DataParallelModel/DataParallelCriterion的內部情況：

下面有兩個特殊情況，並給出瞭解決辦法：

1. 模型輸出了一些張量：你可以利用output_1,output_2 = zip(*predictions)分解它們。
2. 若你不想平行計算損失函式，則可以利用gathered_prdictions = parallel.gather(predictions)收集張量。

分散式訓練

PyTorch中的DistributedDataParallel可以幫助我們在遇到大批量訓練問題時，擁有控制多個伺服器的運算能力。

但值得注意的是：由於對每個節點都要啟動一個獨立的Python訓練指令碼，在設定時需要注意改變工作流程。

每個指令碼在訓練中都會擁有：

1. 它自己的優化器，在每次迭代中都執行一個完整的優化，不需要引數傳輸。
2. 一個獨立的Python直譯器：能夠避免 GIL-freeze

在後面我們將通過程式碼進行討論：

torch.distributed 包能夠為同步分散式運算提供低階原語，基於此構建得到DistributedDataParallel。你可以通過閱讀 文件 以及 教程 對其進行進一步理解。

接下來，我們將使用具有兩個4-GPU的伺服器。

The main server (server 1) has an accessible IP and an open port for communication.

升級Python指令碼以適用分散式訓練

首先，我們需要對指令碼進行升級，使其能夠獨立的在機器（節點）中執行。我們想要完全實現分散式，並且在每個結點的每個GPU上獨立執行程序，這一共需要8個程序。

接下來，初始化分散式後端，封裝模型以及準備資料，這些資料用於在獨立的資料子集中訓練程序。更新後的程式碼如下：

from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import DataLoader

# 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)
for inputs, labels in dataloader:
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地址（示例中為：196.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_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of our training script)

同樣在第二臺機器中執行指令碼：

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_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of our training script)

除了—node_rank引數之外，上述兩個命令相同。

擴充套件

如果你覺得在計算機叢集上執行一組幾乎相同的命令有些枯燥，可點選此處瞭解 GNU並行 。

以上為譯文

本文由阿里云云棲社群組織翻譯。

文章原標題《Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups》，作者：

Thomas Wolf，譯者：Elaine，審校：袁虎。

文章為簡譯，更為詳細的內容，請檢視 原文