隨著深度學習的多項進步，複雜的網絡（例如大型transformer 網絡，更廣更深的Resnet等）已經發展起來，從而需要了更大的內存空間。 經常，在訓練這些網絡時，深度學習從業人員需要使用多個GPU來有效地訓練它們。 在本文中，我將向您介紹如何使用PyTorch在GPU集群上設置分布式神經網絡訓練。
通常，分布式訓練會在有一下兩種情況。
在GPU之間拆分模型：如果模型太大而無法容納在單個GPU的內存中，則需要在不同GPU之間拆分模型的各個部分。跨GPU進行批量拆分數據。當mini-batch太大而無法容納在單個GPU的內存中時，您需要將mini-batch拆分到不同的GPU上。跨GPU的模型拆分
跨GPU拆分模型非常簡單，不需要太多代碼更改。 在設置網絡本身時，可以將模型的某些部分移至特定的GPU。 之後，在通過網絡轉發數據時，數據也需要移動到相應的GPU。 下面是執行相同操作的PyTorch代碼段。
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之間同步變量。
torch.distributed.gather（inputtensor，collectlist，dst）：從所有設備收集指定的inputtensor並將它們放置在collectlist中的dst設備上。torch.distributed.allgather（tensorlist，inputtensor）：從所有設備收集指定的inputtensor並將其放置在所有設備上的tensor_list變量中。torch.distributed.reduce（inputtensor，dst，reduceop = ReduceOp.SUM）：收集所有設備的input_tensor並使用指定的reduce操作（例如求和，均值等）進行縮減。最終結果放置在dst設備上。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翻譯組