DistributedDataParallel介绍
为什么使用DistributedDataParallel
在并行训练中使用DataParallel进行分布式运算时容易遇到list没有复制到GPU等各种问题,导致无法进行分布式计算,然而pytorch推荐的DistributedDataParallel能大大降低bug的发生,因此本文介绍一下DistributedDataParallel函数的正确使用。
DistributedDataParallel使用方法
parser.add_argument('--local_rank', default=-1,type=int)
args= parser.parse_args()
首先需要定义local_rank参数,其表示进程对应的GPU编号,不加的话会报错!
初始化环境
if args.local_rank!=-1:
torch.cuda.set_device(args.local_rank)
device=torch.device("cuda", args.local_rank)
torch.distributed.init_process_group(backend="nccl", init_method='env://')
将模型放入GPU,并进行分布式训练
model.to(device)# 将模型放入GPU中
model= nn.parallel.DistributedDataParallel(teacher_model, device_ids=[args.local_rank],
output_device=args.local_rank)
调整DataLoder函数
这里需要将pin_memory设置为True,可以提高模型训练速度。
train_dataloader= DataLoader(train_datasets, sampler=train_sampler, batch_size=args.train_batch_size,
num_workers=args.num_workers, pin_memory=True)
分布式程序在终端中的启动指令
python-m torch.distributed.launch--nproc_per_node=NUM_GPU train.py
此处NUM_GPU为使用的GPU数量。同时也可以在前面加入CUDA_VISIBLE_DEVICES指定使用GPU的编号。
CUDA_VISIBLE_DEVICES=0,1,2 python-m torch.distributed.launch--nproc_per_node=3 train.py
这是表示使用GPU编号为0,1,2的三块显卡进行训练。0