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

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

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

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

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

　　當訓練批量甚至單個訓練樣本大于GPU内存時，如何訓練模型；如何高效地利用多GPU機器；如何在分布式設備上簡單的使用多個機器。 在一個或多個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：htt第四個步驟（見右上）中，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的内部情況：

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

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

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

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

　　每個腳本在訓練中都會擁有：

　　它自己的優化器，在每次叠代中都執行一個完整的優化，不需要參數傳輸。一個獨立的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，審校：袁虎。

　　,

更多精彩资讯请关注tft每日頭條，我们将持续为您更新最新资讯!