[发明专利]一种神经网络剪枝方法、装置、计算机设备及存储介质

说明书

本申请实施例公开了一种神经网络剪枝方法、装置、计算机设备及存储介质，可以对运行在人工智能处理器上的算子对应的计算图进行剪枝，提高异构系统并行执行神经网络计算任务的效率。

技术领域

本申请涉及信息处理技术领域，尤其涉及一种神经网络剪枝方法、装置、计算机设备及存储介质。

背景技术

现有技术中，在执行一个运行在异构系统(其中，异构系统包括人工智能处理器和通用处理器)上的神经网络计算任务时，通用处理器会将在人工智能处理器上执行的子图里的大量算子编译优化成一个内核函数。具体地，该内核函数作为人工智能处理器上执行的算子与通用处理器上执行的算子一起参与调度。此时，如果神经网络计算任务对应的计算图中的子图划分不合理，造成运行在人工智能处理器上的子图与运行在异构系统中其他设备上的子图之间产生依赖关系，这将影响调度执行阶段异构系统并行执行的效率。那么，如何提高异构系统并行执行神经网络计算任务的效率是亟需解决的技术问题。

发明内容

本申请实施例提供一种神经网络剪枝方法、装置、计算机设备及存储介质，可以避免运行在人工智能处理器上的子图与运行在通用处理器上的子图之间产生依赖关系，可以提高异构系统并行执行神经网络计算任务的效率。

为实现上述目的，第一方面，本申请实施例提供了一种神经网络剪枝方法，所述方法包括：

获取神经网络模型对应的原始计算图，从所述原始计算图中确定第一类算子；其中，第一类算子对应的操作指令能够运行在人工智能处理器上；

根据所述原始计算图中算子之间具有的有向边，获取所述第一类算子构成的计算图，在所述第一类算子构成的计算图中提取原始子图；其中，所述原始子图包含多个输入算子和/或多个输出算子；所述原始子图构成原始子图集合；

对所述原始子图进行剪枝，得到对应的剪枝结果子图；其中，所述剪枝结果子图包含一个输入算子和一个输出算子；

获取所述剪枝结果子图对应的操作指令，以分配至对应的人工智能处理器上执行任务。

为实现上述目的，第二方面，本申请实施例提供了一种神经网络剪枝装置，所述剪枝装置包括：

获取单元，用于获取神经网络模型对应的原始计算图，从所述原始计算图中确定第一类算子；其中，第一类算子对应的操作指令能够运行在人工智能处理器上；

提取原始子图单元，用于根据所述原始计算图中算子之间具有的有向边，获取所述第一类算子构成的计算图，在所述第一类算子构成的计算图中提取原始子图；其中，所述原始子图包含多个输入算子和/或多个输出算子；所述原始子图构成原始子图集合；

剪枝单元，用于对所述原始子图进行剪枝，得到对应的剪枝结果子图；其中，所述剪枝结果子图包含一个输入算子和一个输出算子；

执行单元，用于获取所述剪枝结果子图对应的操作指令，以分配至对应的人工智能处理器上执行任务。

为实现上述目的，第三方面，本申请实施例提供了一种计算机设备，包括处理器和存储器，所述处理器和存储器相互连接，其中，所述存储器用于存储支持计算机设备执行上述方法的计算机程序，所述计算机程序包括程序指令，所述处理器被配置用于调用所述程序指令，执行上述第一方面的方法。

为实现上述目的，第四方面，本申请实施例提供了一种计算机可读存储介质，所述计算机存储介质存储有计算机程序，所述计算机程序包括程序指令，所述程序指令当被处理器执行时使所述处理器执行上述第一方面的方法。

为实现上述目的，第五方面，本申请实施例提供了一种计算机程序，所述计算机程序包括程序指令，所述程序指令当被处理器执行时使所述处理器执行上述第一方面的方法。