[发明专利]张量切分模式的确定方法、装置、计算机设备及介质

说明书

本发明实施例公开了一种张量切分模式的确定方法、装置、计算机设备及介质。该方法包括获取包括多个计算算子的待张量切分的计算图；根据与计算算子对应的可选张量切分模式集，确定与计算图对应的包括多个目标计算算子和目标计算算子的目标张量切分模式的数据条目集；根据各目标张量切分模式，计算各数据条目分别对应的计算耗时；根据在数据条目集中选取的多个数据条目对，生成多个交叉数据条目，并根据各所述交叉数据条目的计算耗时，更新数据条目集；在更新后的数据条目集中获取目标数据条目，从而确定计算图中各计算算子的张量切分模式。本发明实施例的技术方案提供一种可适用复杂场景，能够高效确定算子较优切分模式的方法。

技术领域

本发明实施例涉及人工智能芯片技术，尤其涉及一种张量切分模式的确定方法、装置、计算机设备及介质。

背景技术

现有的人工智能处理器通常在片内集成多个或多组并行计算单元以实现数据的并行化处理，减少神经网络训练或者推理处理的时间，因此芯片对上层的神经网络图编译器有需求。其中一种处理方法是将计算图中所有算子的输入张量数据进行切分，分配给多个计算单元上执行。

上述处理方法面临的一些具体问题如下：在部分场景下不切分可能比切分更划算；整个计算图上算子成千上万，每个算子支持的张量切分形式也通常有多种，因此需要处理的排列组合方式非常多，从中较为快速找到一个最优的或者较优的切分方式是各巨大挑战；网络千变万化，不可能为每个网络进行单独的适配，通用性和可扩展性也是必须要解决的一个问题。

当前主流的计算框架如Pytorch和TensorFlow等不支持上述场景的优化处理。因此，急需探求一种可适用复杂场景的，高效确定算子较优切分模式的方法，满足各种应用场景下的张量切分需求。

发明内容

本发明实施例提供一种张量切分模式的确定方法、装置、计算机设备及介质，以提供一种可适用复杂场景，能够高效确定算子较优切分模式的方法。

第一方面，本发明实施例提供了一种张量切分模式的确定方法，该方法包括：

获取待张量切分的计算图，计算图中包括多个计算算子；

根据与计算算子对应的可选张量切分模式集，确定与计算图对应的数据条目集，数据条目中包括多个目标计算算子和目标计算算子的目标张量切分模式；

根据各所述数据条目中包括的各所述目标张量切分模式，计算与每个数据条目分别对应的计算耗时；

根据在数据条目集中选取的多个数据条目对，生成多个交叉数据条目，并根据各所述交叉数据条目的计算耗时，更新所述数据条目集；

在更新后的数据条目集中获取目标数据条目，并根据目标数据条目，确定计算图中各计算算子的张量切分模式。

第二方面，本发明实施例还提供了一种张量切分模式的确定装置，该装置包括：

计算图获取模块，用于获取待张量切分的计算图，计算图中包括多个计算算子；

数据条目集确定模块，用于根据与计算算子对应的可选张量切分模式集，确定与计算图对应的数据条目集，数据条目中包括多个目标计算算子和目标计算算子的目标张量切分模式；

计算耗时确定模块，用于根据各所述数据条目中包括的各所述目标张量切分模式，计算与每个数据条目分别对应的计算耗时；

数据条目集更新模块，用于根据在数据条目集中选取的多个数据条目对，生成多个交叉数据条目，并根据各所述交叉数据条目的计算耗时，更新所述数据条目集；

张量切分模式确定模块，用于在更新后的数据条目集中获取目标数据条目，并根据目标数据条目，确定计算图中各计算算子的张量切分模式。

第三方面，本发明实施例还提供了一种计算机设备，该设备包括：存储器、处理器及存储在存储器上并可在处理器上运行的计算机程序，所述处理器执行所述程序时实现如本发明任意实施例所述的张量切分模式的确定方法。