MATRIX——专为AI服务的全球分布式资源共享网络 (2/4)

2024/04/09

MANTA的技术原理

MANTA – MATRIX AI Network Training Assistant，这是基于MATRIX高性能区块链打造的分布式自动机器学习平台。这个平台是一个基于分布式技术加速的自动机器学习（AutoML）应用及其部署系统，最初版本的应用领域为计算机视觉中的图像分类以及3D深度估计任务，通过AutoML网络搜索算法搜索一个高精度、低延迟的深度模型，并利用分布式计算技术进行搜索加速。MANTA的核心功能主要分为两个部分：自动机器学习-AutoML与分布式机器学习。

自动机器学习 - 网络结构搜索

MATRIX的自动机器学习采用的网络搜索算法是Once-for-all (OFA)，主要思想是构建一个具备超大网络结构搜索空间并且能衍生出大量子网络的超网（Super Net），然后通过渐进式的剪枝搜索策略，调整超网的网络参数，使得可以从超网中采样不同规模但精度可靠的子网络结构，且无需对采样的子网络进行微调，减小网络部署成本。

以图像分类为例，MANTA用于深度估计的网络结构如图所示。该网络结构拥有5个可搜索单元，每个单元的第一个卷积层的stride都设为2，其余卷积层皆设为1，相当于每一个SU一边进行降维一边提取多通道特征。网络最后有一个SoftMax回归层，通过计算每一类物体的概率，对给定图像进行分类。

图：以图像分类为例的网络结构模型

同时，MANTA用于深度估计的网络结构如下图所示。该网络结构拥有6个可搜索单元，前三个SU主要用来降维和提取左视图和右视图的多尺度特征图，中间有一个Correlation Layer，主要用来计算左右特征图的视差代价，后三个SU进一步对视察代价进行特征提取。随后，网络通过6个转置卷积层，将降维的特征图逐渐放大，并通过U型链接与前面6个SU输出的特征图进行链接，预测不同尺度的深度图。

图：深度估计网络结构模型

分布式机器学习 - 分布式数据/模型并行算法

在MANTA中，将利用GPU分布式并行计算进一步加速网络搜索速度和训练速度。MANTA的分布式机器学习本质上是分布式并行算法，在MANTA中并行算法支持数据并行和模型并行两种加速算法。分布式数据并行算法的主要思想是将每一次迭代中的批量数据分配到不同的GPU中进行前向和后向计算，在每次迭代中，每一个GPU采样的都是相同的子模型。分布式模型并行算法与数据并行不同，在每一次迭代中，MANTA允许不同节点采样不同的子网络，并在各自完成梯度聚合后，进行梯度信息的全局交换。

分布式数据并行算法的流程

１. 从超网中采样4个子网络。

２. 对于每一个子网络，每个节点中的每个GPU都执行下列操作：

（1）读取训练数据，数目为设定的Batch Size；

（2）对当前的子网络进行前向和后向传播操作；

（3）累积后向传播中的梯度信息。

３. 对所有GPU中累积的梯度进行聚合平均，并进行广播（All Reduce操作）。

４. 每个GPU用所得到的全局梯度信息更新自己的超网。

５. 回到步骤1。

分布式模型并行算法（以4个节点为例）的流程

１. 位于同一个节点上的GPU采样同一个子网络模型，总共有4个不同的子网络模型

２. 位于同一个节点上的每个GPU都执行下列操作：

（1）读取训练数据，数目为设定的Batch Size；

（2）对节点上采样的子网络进行前向和后向传播操作；

（3）计算后向传播中的梯度信息；

（4）相同节点内的GPU进行梯度聚合平均，存储于节点的CPU内存。

３. 4个节点将各自在步骤2得到的梯度信息进行聚合平均，并广播到所有GPU。

４. 每个GPU用所得到的全局梯度信息更新自己的超网。

５. 回到步骤1。

图：MANTA的任务训练监控界面