手机最快现场开奖直播
PyTorch Parallel Training
发布日期:2020-10-05 20:13   来源:未知   阅读:

1 为什么不用nn.DataParallel ?1.1 最简单的并行方式

我们在训练时最常用的并行方式就是nn.DataParallel 了, 可以帮助我们(使用单进程控)将模型和数据加载到多个 GPU 中,控制数据在 GPU 之间的流动,协同不同 GPU 上的模型进行并行训练。

只需一行代码,就可以使用多卡进行训练,其中device_ids 用于指定使用的GPU, output_device 用于指定汇总梯度的GPU是哪个:

model = nn.DataParallel(model.cuda(), device_ids=gpus, output_device=gpus[0])

训练模板:# main.pyimport torchimport torch.distributed as distgpus = [0, 1, 2, 3]torch.cuda.set_device('cuda:{}'.format(gpus[0]))train_dataset = ...train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=...)model = ...model = nn.DataParallel(model.to(device), device_ids=gpus, output_device=gpus[0])optimizer = optim.SGD(model.parameters())for epoch in range(100): for batch_idx, (data, target) in enumerate(train_loader): images = images.cuda(non_blocking=True) target = target.cuda(non_blocking=True) ... output = model(images) loss = criterion(output, target) ... optimizer.zero_grad() loss.backward() optimizer.step()1.2 nn.DataParallel 的缺点

在每个训练批次(batch)中,因为模型的权重都是在 一个进程上先算出来 然后再把他们分发到每个GPU上,所以网络通信就成为了一个瓶颈,而GPU使用率也通常很低。2 多进程的 torch.distributed2.1 介绍

在 1.0 之后,官方终于对分布式的常用方法进行了封装,支持 all-reduce,broadcast,send 和 receive 等等。通过 MPI 实现 CPU 通信,通过 NCCL 实现 GPU 通信。官方也曾经提到用DistributedDataParallel 解决 DataParallel 速度慢,GPU 负载不均衡的问题,目前已经很成熟了。

Power by DedeCms