当代研究生应当掌握的并行训练方法(单机多卡)
发布网友
发布时间：2024-10-24 00:45
共1个回答
热心网友
时间：2天前
周五学习时光，闲来无事，决定利用四张 V100-PICE、V100 和 K80 显卡，通过实验比较不同分布式学习库在 ImageNet 上的运行效率，充实一下闲置资源（小小机智一下）。我用 PyTorch 实现了几个示例，包括:
torch.distributed：通过 torch.distributed 提升并行训练性能
torch.multiprocessing：用 torch.multiprocessing 替代传统启动器
apex：借助 apex 提供的混合精度训练和分布式加速
horovod：Uber 开源工具 horovod 的高效应用
基础的分布式 evaluation：利用 all_reduce 和 barrier 等 API
实验结果显示，Apex 的加速效果最为显著，但与 horovod 和 torch.distributed 的差异不大，一般情况下推荐使用内置的 DistributedDataParallel。Dataparallel 的性能较差，不建议使用。接下来会补充 V100/K80 的测试数据。
分享一些关键代码片段：
DataParallel
DataParallel 通过单进程控制多GPU，模型和数据在 GPU 间自动流转，简化了并行训练过程。关键代码如下：
model = DataParallel(model, device_ids=gpus, output_device=gpus[0])
torch.distributed
torch.distributed 提供了分布式训练的高级API，通过 init_process_group 设置通信方式。例如：
init_process_group(backend='nccl', rank=rank, world_size=num_processes)
torch.multiprocessing
使用 torch.multiprocessing spawn 自动创建进程，代码示例：
spawn(main_worker, nprocs=4, args=(local_rank, myargs))
Apex
Apex 优化了混合精度和分布式训练，使用 DistributedDataParallel 时，需要对 loss 进行缩放：
loss = amp.scale_loss(loss, optimizer)
Horovod
Horovod 中，广播模型参数和同步梯度采用:
broadcast_parameters(model.state_dict(), root_rank=0)
完整代码示例请参考 Github。
这只是一个初步的指南，对于工业级实践，欢迎提出建议或贡献。
