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

ゞ浴缸里的玫瑰 2022-10-08 14:30 20阅读 0赞

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

![2021062210093475.png][]

![20210622100750832.png][]

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

![图片][6a139c8874dfc3b8c0b2aa79edc7f5ec.png]

![20210622100810731.png][]

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

![图片][a1285d3aba5eeec1d5ea7253b84f214c.png]

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

![图片][e5a6e3df36e9510658b1a1f48d62162d.png]

如上图所示：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分布式训练时就如上述步骤一样 ，特别简单！

![20210622100830845.png][]

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

（1）导入依赖模块![图片][1eaf58437a8b64f1b383057822ed6d8d.png]       （2）指定batch\_size大小

![图片][7a94002251b50f3a312d69578f44e9e7.png]

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

![图片][ca6d91586a717ae1f049c6769788896e.png]

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

![图片][f85c8f33f8462dd78505602c482cac0b.png]

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

![图片][5f71c5a726766b2c2f46b636de1aadfb.png]

![图片][b473fa43e775fd5f4302739810867c7b.png]

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

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

![图片][ed09d5b921478d904ff1041eceace074.png]

关系图如下：

![图片][2ee326e3a37e29594edab1e6abaa031e.png]

![20210622101608956.png][]

![图片][bf67c0895fcdf1da3521340877b8525a.png]

从源码中可以看出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 数的启动逻辑判断**

![图片][c6c2f8dfd9c83bf5a48c1e836a558f23.png]

从上图中我们可以看到，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()进行分析，其他大致相同。

![图片][78f22963496675faf0b44c339762b48d.png]

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

![20210622101541325.png][]

![图片][11d1da5446b24f4712be0d0a926e5d32.png]

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

![图片][af1a2dfd6bba699475e101f3560472af.png]

![图片][63213f27e87472a9174af983bc31122b.png]

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

![20210622101507682.png][]

**（1）TFMLDataset类介绍**

![图片][a98ae3b0f08cc352590b5bd7d3549125.png]

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

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

![图片][c1786182cd3c5e8dca0ef825ac697ab1.png]

![图片][4d3806a89a1a0c08535af890b5206388.png]

**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。

[3]: https://mp.weixin.qq.com/s?__biz=MzU4NjY4NDU0MQ==&mid=2247485451&idx=1&sn=b3053f0dd9ba8038a0bff3f2e9d7b89e&chksm=fdf6c919ca81400ff05eb2ab1ea878818a1736e166a0d0adae8d7080843f86f4b2a56b77ba71&token=1077402775&lang=zh_CN#rd
[2021062210093475.png]: /images/20221005/cfc397a57cfd41dfa44705968c427b03.png
[20210622100750832.png]: /images/20221005/afd6b1acb54847c29869f416cf81ffe0.png
[6a139c8874dfc3b8c0b2aa79edc7f5ec.png]: /images/20221005/b5375c08fb2542c2bc4befea223af087.png
[20210622100810731.png]: /images/20221005/6a436ce934bb47f39cb73f05df1c191f.png
[a1285d3aba5eeec1d5ea7253b84f214c.png]: /images/20221005/95662f4cf0d74285b168661b5ed8ccd6.png
[e5a6e3df36e9510658b1a1f48d62162d.png]: /images/20221005/ed6f97e1eb3743c0a9b6ba24a3a1ae15.png
[20210622100830845.png]: /images/20221005/ac46357580684e85a619d20d74f51372.png
[1eaf58437a8b64f1b383057822ed6d8d.png]: /images/20221005/4609afb433b343e596efb3d32021f9a7.png
[7a94002251b50f3a312d69578f44e9e7.png]: /images/20221005/9c424b7ee3914b38a8cc5400f8548e2a.png
[ca6d91586a717ae1f049c6769788896e.png]: /images/20221005/dfb9e23397b9404c9b6a7d4078eb0f93.png
[f85c8f33f8462dd78505602c482cac0b.png]: /images/20221005/e37771929c3149d3b83a7091e89dc62a.png
[5f71c5a726766b2c2f46b636de1aadfb.png]: /images/20221005/eb7e8c37bf714a7c84936d06c0035f01.png
[b473fa43e775fd5f4302739810867c7b.png]: /images/20221005/ba3117f859fb48b4ad3d4aa3452c1257.png
[ed09d5b921478d904ff1041eceace074.png]: /images/20221005/3a23b10f71074a43b06f0a2af2d58b10.png
[2ee326e3a37e29594edab1e6abaa031e.png]: /images/20221005/f54e398224be4a2eb0543c341f13d153.png
[20210622101608956.png]: /images/20221005/161f8752985843a1b4ff96320c6ce1a4.png
[bf67c0895fcdf1da3521340877b8525a.png]: /images/20221005/a1ffbf937d2b42d0bb13bd76cadebec0.png
[c6c2f8dfd9c83bf5a48c1e836a558f23.png]: /images/20221005/16cab72a813d41c099f6121d6a2075d1.png
[78f22963496675faf0b44c339762b48d.png]: /images/20221005/5f3009900d7444a79e8542e80cdbdfde.png
[ObjectRef]: http://mp.weixin.qq.com/s?__biz=MzU4NjY4NDU0MQ==&mid=2247484662&idx=1&sn=1d33797c0db369b815af981b64d76850&chksm=fdf6c5e4ca814cf2baaed6449bd33e4aca56e1641f4ec77831f1bf6d31ff1a4f663dc089935f&scene=21#wechat_redirect
[20210622101541325.png]: /images/20221005/6ae2ac44b0204905bc2f2e3a4afd9d83.png
[11d1da5446b24f4712be0d0a926e5d32.png]: /images/20221005/b135ef6803b4450fac9e772496396436.png
[af1a2dfd6bba699475e101f3560472af.png]: /images/20221005/009fd36360c94a379109f4c4c7121bd3.png
[63213f27e87472a9174af983bc31122b.png]: /images/20221005/ab8c5d7645484865bec70b42f6a392ea.png
[20210622101507682.png]: /images/20221005/8db71c8bfd8242db8ebf39747df09141.png
[a98ae3b0f08cc352590b5bd7d3549125.png]: /images/20221005/f04aefb623cc4dc0b84a4f63f59387d4.png
[c1786182cd3c5e8dca0ef825ac697ab1.png]: /images/20221005/3525eaaad6f541a6ae073ad4105c68db.png
[4d3806a89a1a0c08535af890b5206388.png]: /images/20221005/2f31da6692fd4a5d8eecadeee1565be4.png