在pytorch中的多GPU訓練一般有2種DataParallel和DistributedDataParallel ，DataParallel是最簡單的的單機多卡實作，但是它使用多線程模型，并不能夠在多機多卡的環境下使用，是以本文将介紹DistributedDataParallel，DDP 基于使用多程序而不是使用多線程的 DP，并且存在 GIL 争用問題，并且可以擴充到多機多卡的環境，是以他是分布式多GPU訓練的首選。
這裡使用的版本為：python 3.8、pytorch 1.11、CUDA 11.4
如上圖所示，每個 GPU 将複制模型并根據可用 GPU 的數量配置設定資料樣本的子集。
對于 100 個資料集和 4 個 GPU，每個 GPU 每次疊代将處理 25 個資料集。
DDP 上的同步發生在構造函數、正向傳播和反向傳播上。 在反向傳播中梯度的平均值被傳播到每個 GPU。
有關其他同步詳細資訊，請檢視使用 PyTorch 官方文檔：Writing Distributed Applications with PyTorch。
Forking的過程
為了Forking多個程序，我們使用了 Torch 多現成處理架構。一旦産生了程序，第一個參數就是程序的索引，通常稱為rank。
在下面的示例中，調用該方法的所有衍生程序都将具有從 0 到 3 的rank值。我們可以使用它來識别各個程序，pytorch會将rank = 0 的程序作為基本程序。
import torch.multiprocessing as mp
// number of GPUs equal to number of processes
world_size = torch.cuda.device_count()
mp.spawn(<selfcontainedmethodforeachproc>, nprocs=world_size, args=(args,))
GPU 程序配置設定
将 GPU 配置設定給為訓練生成的每個程序。
import torch
import torch.distributed as dist
def train(self, rank, args):
current_gpu_index = rank
torch.cuda.set_device(current_gpu_index)
dist.init_process_group(
backend='nccl', world_size=args.world_size,
rank=current_gpu_index,
init_method='env://'
)
多程序的Dataloader
對于處理圖像，我們将使用标準的ImageFolder加載器，它需要以下格式的樣例資料。
<basedir>/testset/<categoryname>/<listofimages>
<basedir>/valset/<categoryname>/<listofimages>
<basedir>/trainset/<categoryname>/<listofimages>
下面我們配置Dataloader：
from torchvision.datasets import ImageFolder
train_dataset = ImageFolder(root=os.path.join(<basedir>, "trainset"), transform=train_transform)
當DistributedSample與DDP一起使用時，他會為每個程序/GPU提供一個子集。
from torch.utils.data import DistributedSampler
dist_train_samples = DistributedSampler(dataset=train_dataset, num_replicas =4, rank=rank, seed=17)
DistributedSampler與DataLoader進行整合
from torch.utils.data import DataLoader
train_loader = DataLoader(
train_dataset,
batch_size=self.BATCH_SIZE,
num_workers=4,
sampler=dist_train_samples,
pin_memory=True,
)
模型初始化
對于多卡訓練在初始化模型後，還要将其配置設定給每個GPU。
from torch.nn.parallel import DistributedDataParallel as DDP
from torchvision import models as models
model = models.resnet34(pretrained=True)
loss_fn = nn.CrossEntropyLoss()
model.cuda(current_gpu_index)
model = DDP(model)
loss_fn.cuda(current_gpu_index)
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.module.parameters()), lr=1e-3)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7)
訓練
訓練開始時需要在DistributedSampler上設定 epoch，這樣資料在 epoch 之間進行打亂，并且保證在每個 epoch 中使用相同的排序。
for epoch in range(1, self.EPOCHS+1):
dist_train_samples.set_epoch(epoch)
對于DataLoader中的每個批次，将輸入傳遞給GPU并計算梯度。
for cur_iter_data in (loaders["train"]):
inputs, labels = cur_iter_data
inputs, labels = inputs.cuda(current_gpu_index, non_blocking=True),labels.cuda(current_gpu_index, non_blocking=True)
optimizer.zero_grad(set_to_none=True)
with torch.set_grad_enabled(phase == 'train'):
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
loss = loss_fn(outputs, labels)
loss.backward()
optimizer.step()
scheduler.step()
對比訓練輪次的精度，如果更好則存儲模型的權重。
if rank % args.n_gpus == 0 :
torch.save(model.module.state_dict(), os.path.join(os.getcwd(), "scripts/model", args.model_file_name))
在訓練結束時把模型權重儲存在' pth '檔案中，這樣可以将該檔案加載到CPU或GPU上進行推理。
推理
從檔案加載模型:
load_path = os.path.join(os.getcwd(), "scripts/model", args.model_file_name)
model_image_classifier = ImageClassifier()
model_image_classifier.load_state_dict(
torch.load(load_path), strict=False
)
model_image_classifier.cuda(current_gpu_index)
model_image_classifier = DDP(model_image_classifier)
model_image_classifier = model_image_classifier.eval()
這樣就可以使用通常的推理過程來使用模型了。
總結
以上就是PyTorch的DistributedDataParallel的基本知識，DistributedDataParallel既可單機多卡又可多機多卡。
DDP在各程序梯度計算完成之後各程序需要将梯度進行彙總平均,然後再由 rank=0 的程序,将其廣播到所有程序,各程序用該梯度來獨立的更新參數。由于DDP各程序中的模型,初始參數一緻 (初始時刻進行一次廣播),而每次用于更新參數的梯度也一緻的,是以各程序的模型參數始終保持一緻。
DP的處理則是梯度彙總到GPU0,反向傳播更新參數,再廣播參數給其他剩餘的GPU。在DP中,全程維護一個 optimizer,對各個GPU上梯度進行彙總，在主卡進行參數更新,之後再将模型參數 廣播到其他GPU。
是以相較于DP, DDP傳輸的資料量更少,是以速度更快,效率更高。并且如果你使用過DP就會發現，在使用時GPU0的占用率始終會比其他GPU要高，也就是說會更忙一點，這就是因為GPU0做了一些額外的工作，是以也會導緻效率變低。是以如果多卡訓練建議使用DDP進行，但是如果模型比較簡單例如2個GPU也不需要多機的情況下，那麼DP的代碼改動是最小的，可以作為臨時方案使用。
作者：Kaustav Mandal