pytorch DistributedDataParallel多卡並行訓練

pytorch DistributedDataParallel多卡並行訓練

Pytorch 中最簡單的並行計算方式是 nn.DataParallel。
DataParallel 使用單進程控制將模型和數據加載到多個 GPU 中，控制數據在 GPU 之間的流動，協同不同 GPU 上的模型進行並行訓練。

但是DataParallel的缺點十分明顯，各卡之間的負載不均衡，主卡的負載過大。運行時間大概是distributed 的四倍。
所以，下面我們介紹使用distributed 的計算方式。

先設定下local_rank，這是很重要的參數

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=0, type=int, help='node rank for distributed training')
args = parser.parse_args()
print(args.local_rank)

 torch.distributed.init_process_group('nccl', init_method='env://')

 torch.cuda.set_device(args.local_rank)
 device = torch.device("cuda", args.local_rank)

 train_dataset = ModelNet_Loader(root, 'train', args)  #自己寫的類
 test_dataset = ModelNet_Loader(root, 'test', args)    #自己寫的類

 train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
 test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset)

 train_loader = DataLoader(train_dataset, num_workers=2,batch_size=args.batch_size, shuffle=False, drop_last=True,  sampler=train_sampler)       
 test_loader = DataLoader(test_dataset, num_workers=2,batch_size=args.test_batch_size, shuffle=False, drop_last=False, sampler = test_sampler)
   
    #Try to load models
 if args.model == 'STN3d':
        model = STN3d()    #自己的model
 print(str(model))
 model.to(device)

 model = torch.nn.parallel.DistributedDataParallel(model,device_ids=[args.local_rank],output_device=args.local_rank)

然後我們運行如下代碼，啓動訓練

CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2  --master_port 29501  main.py

參考

https://zhuanlan.zhihu.com/p/98535650