DistributedDataParallel(簡稱DDP)是PyTorch自帶的分布式訓練框架, 支持多機多卡和單機多卡, 與DataParallel相比起來, DDP實現了真正的多進程分布式訓練.
[原創][深度][PyTorch] DDP系列第一篇：入門教程
當代研究生應當掌握的並行訓練方法（單機多卡）
DDP的原理和細節推薦上述兩篇文章, 本文的主要目的是簡要歸納如何在PyTorch代碼中添加DDP的部分, 實現單機多卡分布式訓練.
Import部分:
import numpy as np
import torch
import random
import argparse
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
在使用DDP訓練的過程中, 代碼需要知道當前進程是在哪一塊GPU上跑的, 這里對應的本地進程序號local_rank(區別於多機多卡時的全局進程序號, 指的是一台機器上的進程序號)是由DDP自動從外部傳入的, 我們使用argparse獲取該參數即可.
parser = argparse.ArgumentParser(description='Network Parser')
args = parser.parse_args()
local_rank = args.local_rank
獲取到local_rank后, 我們可以對模型進行初始化或加載等操作, 注意這里torch.load()要添加map_location參數, 否則可能導致讀取進來的數據全部集中在0卡上. 模型構建完以后, 再將模型轉移到DDP上:
torch.cuda.set_device(local_rank)
model = YourModel()
# 如果需要加載模型
if args.resume_path:
checkpoint = torch.load(args.resume_path, map_location=torch.device("cpu"))
model.load_state_dict(checkpoint["state_dict"])
# 要在模型初始化或加載完后再進行
# SyncBatchNorm不是必選項, 可以將模型中的BatchNorm層轉換為進程之間同步數據的SyncBatchNorm層, 從而緩解Batch size較小時BN效果差的問題
model = nn.SyncBatchNorm.convert_sync_batchnorm(model).cuda()
model = DDP(model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True)
這里有一個小細節, 假如你的model中用到了隨機數種子來保證可復現性, 那么此時我們不能再用固定的常數作為seed, 否則會導致DDP中的所有進程都擁有一樣的seed, 進而生成同態性的數據:
def init_seeds(seed=0, cuda_deterministic=True):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if cuda_deterministic: # slower, more reproducible
cudnn.deterministic = True
cudnn.benchmark = False
else: # faster, less reproducible
cudnn.deterministic = False
cudnn.benchmark = True
def main():
random_seed = 1234
init_seeds(random_seed+local_rank)
model部分處理完以后, 構建optimizer
optimizer = YourOptimizer()
if args.resume_path:
optimizer.load_state_dict(checkpoint["optimizer"])
對於dataloader部分, 為了讓DDP中同時運行的多個進程使用不同數據, 我們需要引入一個專用的sampler. 注意這里無需再指定shuffle=True, 因為sampler會在后續的set_epoch()幫我們打亂數據.
args.batch_size = args.batch_size // torch.cuda.device_count # 這一步是因為我傳入的參數里batch_size代表所有GPU的batch之和, 所以要除以GPU的數量
train_dataset = YourDataset()
train_sampler = DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_sampler, pin_memory=True, drop_last=True)
對於val_loader, 一般不需要使用上述sampler, 只要保留原始的dataloader代碼即可.
模型和數據都准備好了, 接下來只需簡單的操作:
for epoch in range(0, epoch_num):
train_sampler.set_epoch(epoch) # 這一步是為了讓數據shuffle在每一個epoch都正常工作
train()
在DDP訓練的時候, 我個人習慣的流程是print, tensorboard, validate, 模型sava這類操作都只讓其中一個進程完成, 這樣可以避免冗余信息的出現. 注意DDP的時候保存的是model.module.state_dict().
from torch.utils.tensorboard import SummaryWriter
# 保存模型
if dist.get_rank() == 0:
validate()
state_checkpoint = {
'state_dict': model.module.state_dict(),
'optimizer':optimizer.state_dict()}
torch.save(state_checkpoint, model_name)
)
# 打印信息並在tensorboard繪制
def main():
if dist.get_rank() == 0:
writer = SummaryWriter()
else:
writer = None
train(your_args, writer)
def train(your_args, writer):
if dist.get_rank() == 0:
writer.add_scalar('Train/Loss', your_value, global_step=your_step)
print("Your information")
注意這里創建tensorboard.SummaryWriter的進程和后續寫入的進程要統一, 假如所有進程都創建, 只有一個進程寫入的話, 會導致tensorboard不顯示數據, 假如所有進程都創建, 所有進程都寫入的話, 會導致tensorboard上出現許多條線(多個進程同時寫入)
將上述代碼補充到原來的訓練代碼中后, 就可以模型愉快地進行DDP訓練了, 這里nproc_per_node是使用的GPU數量, 我的代碼中batch_size指的是所有GPU上的batch總和, 使用了4張卡, 所以實際上每張GPU上的mini-batch=8
CUDA_VISIBLE_DEVICES="0,1,2,3" nohup python3 -u -m torch.distributed.launch --nproc_per_node=4 train.py --batch_size 32 --your_args "your args here" > log.out 2>&1 &