Pytorch中的分布式神經網絡訓練|pytorch|bat|拆分|調用_網易訂閱

　　隨著深度學習的多項進步，複雜的網絡（例如大型transformer 網絡，更廣更深的Resnet等）已經發展起來，從而需要了更大的內存空間。 經常，在訓練這些網絡時，深度學習從業人員需要使用多個GPU來有效地訓練它們。 在本文中，我將向您介紹如何使用PyTorch在GPU集群上設置分布式神經網絡訓練。

　　

　　通常，分布式訓練會在有一下兩種情況。

　　1. 在GPU之間拆分模型：如果模型太大而無法容納在單個GPU的內存中，則需要在不同GPU之間拆分模型的各個部分。

　　1. 跨GPU進行批量拆分數據。當mini-batch太大而無法容納在單個GPU的內存中時，您需要將mini-batch拆分到不同的GPU上。

跨GPU的模型拆分

　　跨GPU拆分模型非常簡單，不需要太多代碼更改。 在設置網絡本身時，可以將模型的某些部分移至特定的GPU。 之後，在通過網絡轉發數據時，數據也需要移動到相應的GPU。 下面是執行相同操作的PyTorch代碼段。

　　from torch import nn
class Network(nn.Module):
def __init__(self, split_gpus=False):
super().__init__()
self.module1 = ...
self.module2 = ...
self.split_gpus = split_gpus
if split_gpus: #considering only two gpus
self.module1.cuda(0)
self.module2.cuda(1)
def forward(self, x):
if self.split_gpus:
x = x.cuda(0)
x = self.module1(x)
if self.split_gpus:
x = x.cuda(1)
x = self.module2(x)
return x

跨GPU的數據拆分

　　有3種在GPU之間拆分批處理的方法。

　　· 積累梯度

　　· 使用nn.DataParallel

　　· 使用nn.DistributedDataParallel

積累梯度

　　在GPU之間拆分批次的最簡單方法是累積梯度。 假設我們要訓練的批處理大小為256，但是一個GPU內存只能容納32個批處理大小。 我們可以執行8（= 256/32）個梯度下降迭代而無需執行優化步驟，並繼續通過loss.backward（）步驟添加計算出的梯度。 一旦我們累積了256個數據點的梯度，就執行優化步驟，即調用optimizer.step（）。 以下是用於實現累積漸變的PyTorch代碼段。

　　TARGET_BATCH_SIZE, BATCH_FIT_IN_MEMORY = 256, 32
accumulation_steps = int(TARGET_BATCH_SIZE / BATCH_FIT_IN_MEMORY)
network.zero_grad() # Reset gradients tensors
for i, (imgs, labels) in enumerate(dataloader):

preds = network(imgs) # Forward pass
loss = loss_function(preds, 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
optim.step() # Perform an optimizer step
network.zero_grad() # Reset gradients tensors

　　優點： 不需要多個GPU即可進行大批量訓練。 即使使用單個GPU，此方法也可以進行大批量訓練。

　　缺點： 比在多個GPU上並行訓練要花費更多的時間。

使用nn.DataParallel

　　如果您可以訪問多個GPU，則將不同的批處理拆分分配給不同的GPU，在不同的GPU上進行梯度計算，然後累積梯度以執行梯度下降是很有意義的。

　　

　　多GPU下的forward和backward

　　基本上，給定的輸入通過在批處理維度中分塊在GPU之間進行分配。 在前向傳遞中，模型在每個設備上複製，每個副本處理批次的一部分。 在向後傳遞過程中，將每個副本的梯度求和以生成最終的梯度，並將其應用於主gpu（上圖中的GPU-1）以更新模型權重。 在下一次迭代中，主GPU上的更新模型將再次複製到每個GPU設備上。

　　在PyTorch中，只需要一行就可以使用nn.DataParallel進行分布式訓練。 該模型只需要包裝在nn.DataParallel中。

　　model = torch.nn.DataParallel(model)
...
...
loss = ...
loss.backward()

　　優點：並行化多個GPU上的NN訓練，因此與累積梯度相比，它減少了訓練時間。因為代碼更改很少，所以適合快速原型製作。

　　缺點：nn.DataParallel使用單進程多線程方法在不同的GPU上訓練相同的模型。 它將主進程保留在一個GPU上，並在其他GPU上運行不同的線程。 由於python中的線程存在GIL（全局解釋器鎖定）問題，因此這限制了完全並行的分布式訓練設置。

使用DistributedDataParallel

　　與nn.DataParallel不同，DistributedDataParallel在GPU上生成單獨的進程進行多重處理，並利用GPU之間通信實現的完全並行性。但是，設置DistributedDataParallel管道比nn.DataParallel更複雜，需要執行以下步驟（但不一定按此順序）。

　　將模型包裝在torch.nn.Parallel.DistributedDataParallel中。

　　設置數據加載器以使用distributedSampler在所有GPU之間高效地分配樣本。 Pytorch為此提供了torch.utils.data.Distributed.DistributedSampler。設置分布式後端以管理GPU的同步。 torch.distributed.initprocessgroup（backend ='nccl'）。

　　pytorch提供了用於分布式通訊後端（nccl，gloo，mpi，tcp）。根據經驗，一般情況下使用nccl可以通過GPU進行分布式訓練，而使用gloo可以通過CPU進行分布式訓練。在此處了解有關它們的更多信息https://pytorch.org/tutorials/intermediate/dist_tuto.html#advanced-topics

　　在每個GPU上啟動單獨的進程。同樣使用torch.distributed.launch實用程序功能。假設我們在群集節點上有4個GPU，我們希望在這些GPU上用於設置分布式培訓。可以使用以下shell命令來執行此操作。

　　python -m torch.distributed.launch --nproc_per_node=4
--nnodes=1 --node_rank=0
--master_port=1234 train.py

　　在設置啟動腳本時，我們必須在將運行主進程並用於與其他GPU通信的節點上提供一個空閒埠（在這種情況下為1234）。

　　以下是涵蓋所有步驟的完整PyTorch要點。

　　import argparse
import torch
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import DataLoader
#prase the local_rank argument from command line for the current process
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=0, type=int)
args = parser.parse_args()
#setup the distributed backend for managing the distributed training
torch.distributed.init_process_group('nccl')
#Setup the distributed sampler to split the dataset to each GPU.
dist_sampler = DistributedSampler(dataset)
dataloader = DataLoader(dataset, sampler=dist_sampler)
#set the cuda device to a GPU allocated to current process .
device = torch.device('cuda', args.local_rank)
model = model.to(device)
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
output_device=args.local_rank)
#Start training the model normally.
for inputs, labels in dataloader:
inputs = inputs.to(device)
labels = labels.to(device)
preds = model(inputs)
loss = loss_fn(preds, labels)
loss.backward()
optimizer.step()

　　請注意，上述實用程序調用是針對GPU集群上的單個節點的。 此外，如果要使用多節點設置，則必須在選擇啟動實用程序時選擇一個節點作為主節點，並提供master_addr參數，如下所示。 假設我們有2個節點，每個節點有4個GPU，第一個IP位址為「 192.168.1.1」的節點是主節點。 我們必須分別在每個節點上啟動啟動腳本，如下所示。

　　在第一個節點上運行

　　python -m torch.distributed.launch --nproc_per_node=4 --nnodes=1 --node_rank=0--master_addr="192.168.1.1" --master_port=1234 train.py

　　在第二個節點上，運行

　　python -m torch.distributed.launch --nproc_per_node=4 --nnodes=1 --node_rank=1--master_addr="192.168.1.1" --master_port=1234 train.py

　　其他實用程序功能：

　　在評估模型或生成日誌時，需要從所有GPU收集當前批次統計信息，例如損失，準確率等，並將它們在一臺機器上進行整理以進行日誌記錄。 PyTorch提供了以下方法，用於在所有GPU之間同步變量。

　　1. torch.distributed.gather（inputtensor，collectlist，dst）：從所有設備收集指定的inputtensor並將它們放置在collectlist中的dst設備上。

　　1. torch.distributed.allgather（tensorlist，inputtensor）：從所有設備收集指定的inputtensor並將其放置在所有設備上的tensor_list變量中。

　　1. torch.distributed.reduce（inputtensor，dst，reduceop = ReduceOp.SUM）：收集所有設備的input_tensor並使用指定的reduce操作（例如求和，均值等）進行縮減。最終結果放置在dst設備上。

　　1. torch.distributed.allreduce（inputtensor，reduce_op = ReduceOp.SUM）：與reduce操作相同，但最終結果被複製到所有設備。

　　有關參數和方法的更多詳細信息，請閱讀torch.distributed軟體包。 https://pytorch.org/docs/stable/distributed.html

　　例如，以下代碼從所有GPU提取損失值，並將其減少到主設備（cuda：0）。

　　#In continuation with distributedDataParallel.py abovedef get_reduced_loss(loss, dest_device):
loss_tensor = loss.clone()
torch.distributed.reduce(loss_tensor, dst=dest_device)
return loss_tensorif args.local_rank==0:
loss_tensor = get_reduced_loss(loss.detach(), 0)
print(f'Current batch Loss = {loss_tensor.item()}'優點：相同的代碼設置可用於單個GPU，而無需任何代碼更改。 單個GPU設置僅需要具有適當設置的啟動腳本。

　　缺點： BatchNorm之類的層在其計算中使用了整個批次統計信息，因此無法僅使用一部分批次在每個GPU上獨立進行操作。 PyTorch提供SyncBatchNorm作為BatchNorm的替換/包裝模塊，該模塊使用跨GPU劃分的整個批次計算批次統計信息。 請參閱下面的示例代碼以了解SyncBatchNorm的用法。

　　network = .... #some network with BatchNorm layers in itsync_bn_network = nn.SyncBatchNorm.convert_sync_batchnorm(network)
ddp_network = nn.parallel.DistributedDataParallel(
sync_bn_network,
device_ids=[args.local_rank], output_device=args.local_rank)

總結

　　· 要在GPU之間拆分模型，請將模型拆分為submodules，然後將每個submodule推送到單獨的GPU。

　　· 要在GPU上拆分批次，請使用累積梯度nn.DataParallel或nn.DistributedDataParallel。

　　· 為了快速進行原型製作，可以首選nn.DataParallel。

　　· 為了訓練大型模型並利用跨多個GPU的完全並行訓練，應使用nn.DistributedDataParallel。

　　· 在使用nn.DistributedDataParallel時，用nn.SyncBatchNorm替換或包裝nn.BatchNorm層。

　　作者：Nilesh Vijayrania

　　deephub翻譯組

特別聲明：以上內容(如有圖片或視頻亦包括在內)為自媒體平臺「網易號」用戶上傳並發布，本平臺僅提供信息存儲服務。

Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.