分布式计算引擎—Ray SGD 设计原理及核心实现

点击 3分钟秒懂大数据 查看原文链接，关注博主公众号： 3分钟秒懂大数据，带你畅游大数据、算法的知识海洋！

现如今，随着深度模型的增大，只在单节点上进行模型训练，将会花费数天甚至数周的时间，基于此，Ray 社区提出基于 Ray 实现一个分布式训练包装器，也就是 Ray SGD，Ray SGD具备如下特点：

上面介绍到，Ray SGD是一个分布式训练包装器，主要对深度学习框架Tensorflow、Pytorch进行包装，通过num_workers参数 来控制Ray节点的副本，每个副本由Ray Actor进行管理。通过下面这张图可以看出Ray SGD的架构。

由于是对Tensorflow 和Pytorch 进行包装，下面我们来看一下具体的使用规则：

如上图所示：RaySGD对Tensoflow 或者Pytorch进行分布式训练时，包含以下几个步骤：

1、导入依赖模块；

2、由于Ray SGD模块依赖于Ray Core模块，所以在训练时，必须启动Ray,执行图上步骤1.

3、以Pytorch为例，RaySGD在训练时，主要包装了一个TorchTrainer类，在这个包装类里面执行了图上步骤2， 导入数据，创建模型，对模型配置损失函数等，注册模型，加载数据等步骤，然后将其封装成一个类，传给TorchTrainer类。

4、执行图上步骤3，通过num_works 指定副本数量，是否使用GPU，每个batch_size的大小

5、执行trainer.train()方法进行训练。

RaySGD对Tensorflow和Pytorch分布式训练时就如上述步骤一样 ，特别简单！

下面我们通过分布式tensorflow案例来分析源码是如何设计的，案例如下：

（1）导入依赖模块       （2）指定batch_size大小

（3）加载数据，注册训练集和测试集

（4）定义模型，配置损失函数

（5）使用raySGD 对tensorflow的包装类 TFTrainer,进行训练

通过上述案例，我们可以看出，包装类是TFTrainer，将模型，数据、batch_size大小，以及副本数传进去，调用train()方法进行训练即可。所以现在通过源码，我们分析一下，RaySGD是如何封装Tensorflow的。

通过github链接访问Ray-master源码：https://github.com/ray-project/ray，通过定位，分布式RaySGD-Tensorflow源码位置如下：

关系图如下：

从源码中可以看出TFTrainer是一个Tensorflow 训练类，其中定义了许多参数：

（1）model_creater: 接受一个字典类型的配置，返回一个TF 模型；

（2）data_creator: 使用配置创建训练集和验证集；

（3）config：用来定义batch_size的大小，将其传入data_creator，model_creater；

（4）num_replicas: 用于分布式训练 worker的数量；

（5）num_cpus_per_worker: 每个worker请求的cpu个数。
并且根据num_replicas的设置进行许多方法的定义。

（1）worker 数的启动逻辑判断

从上图中我们可以看到，TFTrainer源码对num_replicas进行了逻辑判断，当等于1时，ray远程调用remote，启动1个worker，当不等于1时，根据实际的数量启动相应的worker数，进行分布式计算。

（2）train方法

TFTrainer类中，定义了多个方法，包含train()、validate()、get_model()、save()、restore()、shutdown()、get_model_from_state()等方法，这里以train()进行分析，其他大致相同。

从代码中可以看出，w.step.remote() for w in self.workers 这句是通过设置的worker数量获取到所有的ObjectRef，如不懂objectRef, 可观看之前发布的博文进行了解ObjectRef，然后通过ray.get获取到所有的Object对象，将第一个Object对象进行复制进行返回。得到stats.  其中w.step中的step在TFRunner中讲解。

通过上述源码可以看出，TFRunner用于管理训练的Tensorflow模型，其中传入的参数全部是 上述TFTrainer定义好的变量，并且也定义了许多方法，讲解一下step方法，因为上文用到了。

其中verbose表示输出的训练数据，重点在于self.model.fit, 从这句代码就可以知道，我们在Tensorflow案例中定义的model_create方法最终会调用tensorflow的fit方法对模型进行训练。

（1）TFMLDataset类介绍

通过这个类可以得到许多信息，首先Ray SGD进行训练的数据结构为MLdataset，而TensorFlow 训练的数据结构为TFMLdataset,所以在使用TFTrainer进行训练时，需要将MLdataset转为TFMLdataset。其次，在转换过程中，需要指定特征列和目标列，最后返回一个ds。

（2）获取分片，返回所有的Tensorflow 分片数据

set_num_shards()方法：

在转换成ds后，需要判断ds的分片数量是否和指定的分片数量相同，若不同，重新分区，返回一个指定数量的dataset，如图1。

get_shard()方法

图2中it 表示 所有的分片dataset。

make_generator()方法 在 get_shard()方法里面

图3表示对所有的分片dataset进行遍历，拿到一个分片数据，图4表示 拿到一个分片数据的所有特征列，图5表示拿到一个分片数据的特征列，通过图6进行递归调用make_generator 来返回一个被给的tensorflow 分片数据来供tensorflow进行训练。

上述已将所有的源码进行分析，现在将整个流程进行串通起来：

（1）当执行Tensorflow案例的Trainer.train()方法后，

（2）train()会根据 设置的worker数量去获取所有的stats对象，在获取时，会调用TFRunner类中的step方法；

（3）在step 方法中会调用model.fit方法进行模型训练，传入的参数为train_dataset

（4）由于train_dataset为tensorflow的数据结构，而RaySGD的数据结构为MLDataset，所以需要调用TFMLDataset类将MLDataset转为TFMLDataset，即最后get_shard返回的df。