vivo手机|PyTorch Parallel Training
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 负载不均衡的问题 ， 目前已经很成熟了 。
文章图片
文章图片
与 DataParallel 的单进程控制多 GPU 不同 ， 在 distributed 的帮助下 ， 我们只需要编写一份代码 ， torch 就会自动将其分配给n个进程 ， 分别在n个 GPU 上运行 。
和单进程训练不同的是 ， 多进程训练需要注意以下事项：DistributedSamplerargs.local_rank2.2 使用方式2.2.1 启动方式的改变
在多进程的启动方面 ， 我们不用自己手写 multiprocess 进行一系列复杂的CPU、GPU分配任务 ， PyTorch为我们提供了一个很方便的启动器torch.distributed.lunch 用于启动文件 ， 所以我们运行训练代码的方式就变成了这样：
CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py
其中的--nproc_per_node 参数用于指定为当前主机创建的进程数 ， 由于我们是单机多卡 ， 所以这里node数量为1 ， 所以我们这里设置为所使用的GPU数量即可 。2.2.2 初始化
在启动器为我们启动python脚本后 ， 会通过参数local_rank 来告诉我们当前进程使用的是哪个GPU ， 用于我们在每个进程中指定不同的device：def parse():parser = argparse.ArgumentParser()parser.add_argument('--local_rank', type=int, default=0)args = parser.parse_args()return argsdef main():args = parse()torch.cuda.set_device(args.local_rank)torch.distributed.init_process_group('nccl',init_method='env://')device = torch.device(f'cuda:{args.local_rank}')...
其中 torch.distributed.init_process_group 用于初始化GPU通信方式（NCCL）和参数的获取方式（env代表通过环境变量）2.2.3 DataLoader
推荐阅读
- iQOO手机|畅爽游戏，极速充电！iQOO 5 Pro定义性能旗舰新标准
- 手机芯片|首发A14处理器 iPad Air4现货就绪：苹果将于下周三公布上市安排
- 界面新闻|手机欠费超三个月将上征信，闲置卡以及宽带用户要小心了
- iQOO手机|iQOO 5不只是性能“怪兽”，极致视听同样深得用户芳心
- 产品|一加手机CEO刘作虎：产品线将迎来“新成员”率先海外推出
- 家长|微软推出Xbox家长控制手机应用
- 触目惊心|触目惊心！一部手机丢失后有多可怕？多平台中招，支付宝、银联紧急回应
- 新机发布|买手机的朋友不妨再等等，10月还有这五款机型值得期待！
- 手机|OPPO 液体镜头专利有多厉害？这两点传统镜头根本比不上
- 华为手机|不用等华为Mate40，这四部华为顶级旗舰也不错，可升鸿蒙