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

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

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

1. 跨GPU進行批次拆分資料。當mini-batch太大而無法容納在單個GPU的記憶體中時，您需要將mini-batch拆分到不同的GPU上。

from torch import nnclass 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, 32accumulation_steps = int(TARGET_BATCH_SIZE / BATCH_FIT_IN_MEMORY)network.zero_grad()                            # Reset gradients tensorsfor 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 <OTHER TRAINING ARGS>

在設定啟動指令碼時，我們必須在將執行主程序並用於與其他GPU通訊的節點上提供一個空閒埠（在這種情況下為1234）。

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

import argparseimport torchfrom torch.utils.data.distributed import DistributedSamplerfrom torch.utils.data import DataLoader#prase the local_rank argument from command line for the current processparser = argparse.ArgumentParser()parser.add_argument("--local_rank", default=0, type=int)args = parser.parse_args()#setup the distributed backend for managing the distributed trainingtorch.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 <OTHER TRAINING ARGS>

在第二個節點上，執行

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 <OTHER TRAINING ARGS>

其他實用程式功能：

在評估模型或生成日誌時，需要從所有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層。