PyTorch分布式训练

概览

PyTorch 是一个 Python 优先的深度学习框架，能够在强大的 GPU 加速基础上实现张量和动态神经网络。PyTorch的一大优势就是它的动态图计算特性。

License ：MIT License

官网：http://pytorch.org/

GitHub：https://github.com/pytorch/pytorch

Pytorch 是从Facebook孵化出来的，在0.4的最新版本加入了分布式模式，比较吃惊的是它居然没有采用类似于TF和MxNet的PS-Worker架构。而是采用一个还在Facebook孵化当中的一个叫做gloo的家伙。

PyTorch分布式

官方教程：http://pytorch.org/docs/master/distributed.html

其实这种三种backend对现在我们来说可以说是没得选的，只有gloo支持GPU

这里引入了一个新的函数model = torch.nn.parallel.DistributedDataParallel(model)为的就是支持分布式模式

不同于原来在multiprocessing中的model = torch.nn.DataParallel(model,device_ids=[0,1,2,3]).cuda()函数，这个函数只是实现了在单机上的多GPU训练，根据官方文档的说法，甚至在单机多卡的模式下，新函数表现也会优于这个旧函数。

这里要提到两个问题：

• 每个进程都有自己的Optimizer同时每个迭代中都进行完整的优化步骤，虽然这可能看起来是多余的，但由于梯度已经聚集在一起并跨进程平均，因此对于每个进程都是相同的，这意味着不需要参数广播步骤，从而减少了在节点之间传输张量tensor所花费的时间。
• 另外一个问题是Python解释器的，每个进程都包含一个独立的Python解释器，消除了来自单个Python进程中的多个执行线程，模型副本或GPU的额外解释器开销和“GIL-thrashing”。 这对于大量使用Python运行时的模型尤其重要。

Gloo

项目地址：https://github.com/facebookincubator/gloo

是一个类似MPI的通信库，你不需要考虑内存数据的拷贝，只需要实现逻辑就可以。

初始化

torch.distributed.init_process_group(backend, init_method='env://', **kwargs)

参数说明：

• backend(str): 后端选择，包括上面那几种 tcp mpi gloo
• init_method(str，optional): 用来初始化包的URL我理解是一个用来做并发控制的共享方式
• world_size(int, optional):参与这个工作的进程数
• rank(int,optional): 当前进程的rank
• group_name(str,optional): 用来标记这组进程名的

解释一下init_method()也有这三种方式，具体可参看http://pytorch.org/docs/master/distributed.html

• file:// 共享文件系统（要求所有进程可以访问单个文件系统）有共享文件系统可以选择
• tcp:// IP组播（要求所有进程都在同一个网络中）比较好理解，不过需要手动设置rank
• env:// 环境变量（需要您手动分配等级并知道所有进程可访问节点的地址）默认是这个

MINIST数据集示例

from __future__ import print_function
import</