一种基于蚁群算法的算子融合方法及应用技术

本发明专利技术公开了一种基于蚁群算法的算子融合方法，所述方法包括如下步骤：导入深度学习模型，利用TVM内置函数将其转换为Relay表达式；对Relay表达式进行后序深度优先搜索得到后序深度优先生成树作为DAG；基于DAG利用LCA算法求解每个节点对应所有输出节点的公共祖先节点，从而求得每个节点的后支配点，最终得到后支配树；基于后支配树，先对不含有融合分支的节点进行算子融合，然后对含有融合分支的节点逐一通过蚁群算法求解最佳融合路径，最后再按照最佳融合路径进行算子融合。本发明专利技术可以对TVM中的算子融合算法进行改进，从而在探索融合路径多元化的同时实现算子融合时间效率上的提升。本发明专利技术还公开了上述方法在提升算子融合速度，优化深度学习模型中的应用。优化深度学习模型中的应用。优化深度学习模型中的应用。

【技术实现步骤摘要】
一种基于蚁群算法的算子融合方法及应用


[0001]本专利技术属于深度学习模型优化
，尤其涉及一种基于蚁群算法的算子融合方法及应用。

技术介绍

[0002]蚁群算法(Ant Colony Algorithm)是一种群体智能优化算法，是基于自然界中蚂蚁觅食行为的启发式算法。蚁群算法通过模拟蚂蚁在觅食过程中的信息素沉淀和信息素挥发的行为构建路径优劣的评判标准，通过轮盘赌法选择路径，从而在搜索空间中寻找最优解。蚁群算法适合应用于大规模应用场景，能够高效地在搜索空间中找到最优解。
[0003]面向复杂实际问题的深度学习模型存在计算和访存开销大的问题，算子融合(Operator Fusion)是深度学习模型优化的关键技术。算子融合指的是将深度学习模型中多个存在数据依赖关系的算子融合成一个算子以减少内存访问和数据搬移的次数，从而提高深度学习模型的计算的效率和性能。
[0004]TVM是一个端到端的编译和优化堆栈，能够实现图层次和张量层次的优化并且能够将深度学习的工作负载部署到多种硬件后端上。算子融合是TVM图层次的一种优化方法。
[0005]在TVM中，算子融合是基于四种算子类型和三种融合规则来实现的。其中，四种算子类型分别是Injective(单射性)、Reduction(简约)、Complex
‑
Out
‑
Fusable(逐元素复用映射到输出)和Opaque(不能被复用)，三种融合规则分别是多个Injective算子可以融合为一个Injective算子、作为输入的Injective算子可以与Reduction算子融合，以及Complex
‑
Out
‑
Fusable算子可以将Element
‑
Wise级别的算子融合到其输入中。此外，算子融合算法的执行基于Relay表达式构建的后支配树，其融合路径就是深度遍历DAG所得到的算子序列。
[0006]TVM中的算子融合以遍历顺序作为融合路径，其算子融合的效率为算法的时间复杂度做出了让步。因此，从算子融合路径的角度出发，TVM的算子融合在时间效率上有可提升的空间。

技术实现思路

[0007]为了解决现有技术存在的不足，本专利技术的目的是提供一种基于蚁群算法的算子融合方法，用于对TVM中的算子融合算法进行改进。本专利技术通过在后支配树的每个子树处进行蚁群算法的调度寻找算子融合时间最少的融合路径，从而可以在探索融合路径多元化的同时实现算子融合时间效率上的提升。
[0008]实现本专利技术目的的具体技术方案是：
[0009]一种基于蚁群算法的算子融合方法，该方法包括以下具体步骤：
[0010]步骤1：从Tensorflow、MXNet、ONNX或PyTorch等TVM所支持的深度学习框架中导入深度学习模型，利用TVM中的前端模块将导入的深度学习模型转化为Relay表达式；
[0011]步骤2：遍历Relay表达式结构构建DAG(有向无环图)，具体为：
[0012]对Relay表达式进行后序深度优先搜索，即将Relay表达式的最终输出节点作为根
节点并将每个节点的输入边看作输出边再进行深度优先搜索；
[0013]将Relay表达式后序深度优先搜索的序列倒序存放于列表中，与此同时对Relay表达式进行后序深度优先搜索所得到的后序深度优先生成树即为需要构建的DAG；本专利技术中，所述有向无环图DAG作为求解后支配树的基础，根据DAG中节点的分布以及联系求解后支配树。
[0014]步骤3：基于DAG利用支配树算法构建后支配树，具体为：
[0015]后支配树的生成要基于后序深度优先生成树，遍历步骤2所得的后序深度优先搜索节点列表，对每个节点进行后支配节点的求解从而最终得到后支配树；
[0016]其中，为每个节点求解后支配节点时所使用的是最近公共祖先算法(Lowest Common Ancestor,LCA)；
[0017]所述最近公共祖先算法一般可以用于在树或者有向无环图中查找两个节点的最近公共祖先，在本专利技术中多次使用最近公共祖先算法求解含有多个输入边的节点的后支配节点。步骤4：基于步骤3构建的后支配树进行算子融合，具体为：
[0018]从头遍历步骤2所得的后序深度优先搜索节点列表，首先将所有不含融合分支的节点利用TVM的三种融合规则进行融合，然后对所有含有融合分支的节点逐一通过蚁群算法求解最佳融合路径，最后按照最佳融合路径对含有融合分支的节点基于TVM的三种融合规则进行算子融合操作；本专利技术中，所述融合分支是指当前节点在后支配树的结构中含有的多个分支节点。所述融合规则是指1)多个Injective算子可以融合为一个Injective算子、2)作为输入的Injective算子可以与Reduction算子融合，以及3)Complex
‑
Out
‑
Fusable算子可以将Element
‑
Wise级别的算子融合到其输入中。
[0019]步骤3，对每个节点进行后支配节点的求解从而最终得到后支配树，具体为：
[0020]基于后序深度优先生成树，如果当前节点不存在输入边，则其后支配节点为None，即不存在后支配节点；
[0021]如果当前节点存在一个输入边，则当前节点的输入节点就是其后支配节点；
[0022]如果当前节点存在多个输入边，则对所有的输入节点进行如下操作：
[0023]第一步是利用LCA算法求解前两个输入节点的最近公共祖先；
[0024]第二步是利用LCA算法求解当前最近得到的最近公共祖先和后一个输入节点的最近公共祖先；
[0025]不断重复第二步直到没有后一个输入节点，最终求解所得为所有输入节点的最近公共祖先，即当前节点的后支配节点。
[0026]所有节点的后支配节点求解完成之后，即可得到后支配树。
[0027]当前节点含有多个输入节点，将这些输入节点按读取的顺序排列，假设输入节点为1，2，3，4，先求1，2的LCA，即为lca1，然后再求lca1和3的LCA，依次类推。
[0028]所述LCA算法，具体为：
[0029]对于两个节点，判断两个节点的深度，节点深度是指从根节点到该节点的路径长度；
[0030]将深度较大的节点沿着双亲节点的方向提升，直到和深度较小的节点深度一致；
[0031]将两个节点同时沿着双亲方向提升使得它们变为同一个点；
[0032]最终得到的该同一个点就是要求的两个节点的最近公共祖先。
[0033]步骤4所述对所有含有融合分支的节点逐一通过蚁群算法求解最佳融合路径，具体为：
[0034]初始化参数，将蚂蚁数量设为融合分支的个数，信息素因子alpha设为3，启发函数因子beta设为5，信息素挥发因子r设为0.6，最大迭代次数设为1000，信息素总量设为10
‑6；
[0035]将融合分支节点即待融合算本文档来自技高网...

【技术保护点】

【技术特征摘要】
1.一种基于蚁群算法的算子融合方法，其特征在于，所述方法包括如下步骤：步骤一、从深度学习框架中导入深度学习模型，利用TVM中的前端模块将导入的深度学习模型转化为Relay表达式；步骤二、遍历Relay表达式结构构建有向无环图DAG；步骤三、基于有向无环图DAG利用支配树算法构建后支配树；步骤四、利用步骤三中构建获得的后支配树进行算子融合。2.如权利要求1所述的算子融合方法，其特征在于，步骤一中，所述深度学习模型为TVM支持的深度学习框架，包括Tensorflow、ONNX、MXNet或PyTorch。3.如权利要求1所述的算子融合方法，其特征在于，步骤一中，将深度学习模型转化为Relay表达式的步骤包括：将模型的输入转化为张量tensor并且按照模型的输入要求对张量tensor的尺寸shape进行变化，然后利用TVM内置的函数将深度学习模型转化为Relay表达式的形式。4.如权利要求1所述的算子融合方法，其特征在于，步骤二中，将Relay表达式最终的输出节点看作根节点同时将每个节点的输入边看作输出边再进行深度优先搜索，构建所述有向无环图DAG，进一步包括如下步骤：步骤2.1、判断当前遍历节点的节点类型，所述节点类型包括常量节点ConstantNode、函数调用节点CallNode、函数定义节点FunctionNode；步骤2.2、根据步骤一中判断获得的节点类型选择对应的遍历函数，获取节点信息构建与遍历节点对应的DAG节点；步骤2.3、结合后序深度优先搜索方法重复步骤2.1和步骤2.2，将后序优先搜索序列倒序记录在列表中，对Relay表达式进行后序深度优先搜索所得到的后序深度优先生成树，即有向无环图DAG。5.如权利要求4所述的算子融合方法，其特征在于，步骤三中，根据包括后续优先搜索序列倒序记录的列表中的节点顺序，从后向前遍历，结合DAG和最近公共祖先LCA算法求解每个节点后支配节点，构建后支配树；进一步包括如下步骤：步骤3.1、基于DAG，判断当前遍历节点的输入节点的数量；步骤3.2、如果当前遍历节点只有一个输入节点，那么当前遍历节点的后支配节点就是其输入节点；或，如果当前遍历节点有多个输入节点，那么利用LCA算法两两计算求解出多个输入节点的最近公共祖先节点，则当前遍历节点的后支配节点即为多个输入节点的最近公共祖先节点；步骤3.3、所有节点的后支配节点求解完成后，得到后支配树。6.如权利要求5所述的算子融合方法，其特征在于，所述LCA算法，具体为：对于两个节点，判断两个节点的深度，节点深度是指从根节点到该节点的路径长度；将深度较大的节点沿着双亲节点的方向提升，直到和深度较小的节点深度一致；将两个节点同时沿着双亲方向提升使得它们变为同一个点；最终得到的该同一个点就是要求的两个...

【专利技术属性】
技术研发人员：张越龄，杨雪，王江涛，
申请(专利权)人：华东师范大学，
类型：发明
国别省市：