深度学习框架如何自动选择最快的算法?Fast Run 让你收获最好的性能!

对于深度学习框架来说,网络的训练/推理时间是用户非常看中的。在实际生产条件下,用户设计的 NN 网络是千差万别,即使是同一类数学计算,参数也各不相同。如果没有针对性的优化,框架就完全丧失竞争力。因此,在一类数学计算中,开发者们会开发多种高效的算法,分别适用于不同的参数,以保证网络的性能。接下来开发者们需要解决一个新问题,当计算参数确定以后,如何让最快的算法执行该计算。

大部分框架靠先验的经验选择算法,MegEngine 亦总结有优秀的先验经验值,实现计算时自动选择算法。但是依靠经验不能保证一定选择了最快的算法。很多实际场景中,用户希望网络有最极致的性能。为此,MegEngine 设计了专门的流程,可以为每个计算自动选择最快的算法,从而保证整个网络的运行时间最短。并且同时能够将计算的参数和其对应的算法信息以及设备信息记录到内存或文件,当用户再次运行网络时,可以直接获取性能最好的算法。这一提升性能的流程被称为 Fast Run,它能让 MegEngine 的用户运行不同的网络时都能收获最好的性能。

二、Fast Run 简述

目前,主流的框架几乎都使用了算子(Operator)的概念来抽象数学计算,如卷积算子,矩阵乘算子等。MegEngine 也使用了算子 这一概念。此外,在底层,我们开发了名为 MegDNN 的计算库,用以完成实际的数学计算。 MegDNN 仅提供数学计算能力。MegDNN 的顶层也是按照算子的概念组织的,对不同的后端,分别封装了 MegDNN 算子。一个 MegDNN 算子内部则可能有多个该算子的算法,MegEngine 将算法抽象为 Algorithm,一个 Algorithm 对象可以完成该算子的计算。

以卷积算子为例,ARM 上,MegEngine 实现了非常通用的 Im2col 算法,有特定条件下性能卓越的 Winograd 算法,有在小尺寸卷积时高性能的 Direct 直接卷积算法等。CUDA 上,有调用 cuDNN 库函数的方法等等。从 MegEngine 算子到 MegDNN 算子再到算法的关系如下图所示:

深度学习框架如何自动选择最快的算法?Fast Run 让你收获最好的性能!

一个 MegEngine 算子可能持有一个或多个 MegDNN 算子来完成计算,一个 MegDNN 算子需要从多个算法对象中选择一个来执行计算。为了极致的计算性能,需要在开始网络计算之前,给 MegDNN 算子选好最快的算法。

Fast Run 的思路很直接,在网络计算开始之前,将每个 MegDNN 算子中所有可行的算法全部运行一次(Profiling),并将性能数据记录下来,将最快的算法设置给 MegDNN 算子。Fast Run 成立的前提条件是算法运行时间是稳定的,这样比较每个算法的 Profiling 数据才有意义。

最后是确定 Fast Run 执行的时间点。MegEngine 有统一的内存管理,各 MegEngine 算子需要在计算开始前向内存规划单元申请足够的计算时内存,这一内存包括了其内部的 MegDNN 算子计算时需要的内存,而 MegDNN 算子计算时需要的内存完全由算法决定。这就要求,MegDNN 此刻已经确定了将要使用的算法。自然地,MegEngine 选择在调用该接口之前执行 Fast Run 流程。这样,当 Fast Run 流程完成时,各 MegDNN 算子都设置了性能最好的算法。

Fast Run 执行的代价是显然的,它会显著增加第一次网络执行的时间。Fast Run 的流程如下图:

深度学习框架如何自动选择最快的算法?Fast Run 让你收获最好的性能!

Fast Run 有下面两种使用方式,区别在于上图中写入的 Cache 文件不同:

离线 Fast Run,离线 Fast Run 分两步,分别在不同的进程中完成。第一步先将整个网络计算执行一遍,这一过程中,Fast Run 会将各个算法的性能数据写到一个专门的数据结构中,最后数据被统一写入一个 Cache 文件,随后进程退出,这个过程称之为“搜参”。第二步,加载同样的网络,通过 MegEngine 的接口将 Cache 文件读入。可以看出,离线 Fast Run甚至可以在不同的设备上进行。

在线 Fast Run,在线 Fast Run 在同一个进程完成的。前半段与离线 Fast Run 的流程相同,Fast Run 后,各算法的性能数据保存在内存中的一个数据结构之中。此时,进程不会退出。后续可以给网络加载不同的输入数据,此时各 MegDNN 算子中已设置好性能最好的算法。并且,也可以初始化另外的网络,亦可以像离线 Fast Run 的后半部分一样,从当前的数据结构中读取算法。

总的来说,Fast Run 提供搜参和记录的功能。它的作用是给网络中的各个 MegDNN 算子选择当前参数下性能最好的算法。由于 Fast Run 对每个 MegDNN 算子执行同样的操作,因此它在前向推理和反向传播时都能使用。目前,MegEngine 支持 CUDA、CPU、ROCM 三个后端的 Fast Run ,MegEngine 的用户们在训练和部署时,均广泛使用 Fast Run。

三、Fast Run 原理

Fast Run 中,Profiling 一个 MegDNN 算子并设置算法,会经历 4 个步骤,其流程如下图示:

内容版权声明:除非注明,否则皆为本站原创文章。

转载注明出处:https://www.heiqu.com/zzwjgx.html