动态图

下载Notebook下载样例代码查看源文件

基础能力

在MindSpore中,动态图模式又被称为PyNative模式,为默认模式,也可以通过set_context(mode=PYNATIVE_MODE)来设置成动态图模式。在脚本开发和网络流程调试中,动态图模式下更容易调试,支持执行单算子、普通函数和网络、以及单独求梯度的操作。

在PyNative模式下,用户可以使用完整的Python API,此外针对使用MindSpore提供的API时,框架会根据用户选择的硬件平台(Ascend,GPU,CPU),将算子API的操作在对应的硬件平台上执行,并返回相应的结果。框架整体的执行过程如下:

process

通过前端的Python API,调用到框架层,最终到相应的硬件设备上进行计算。例如:完成一个加法。

[2]:
import numpy as np
import mindspore as ms
import mindspore.ops as ops

ms.set_context(mode=ms.PYNATIVE_MODE, device_target="CPU")
x = ms.Tensor(np.ones([1, 3, 3, 4]).astype(np.float32))
y = ms.Tensor(np.ones([1, 3, 3, 4]).astype(np.float32))
output = ops.add(x, y)
print(output.asnumpy())
[[[[2. 2. 2. 2.]
   [2. 2. 2. 2.]
   [2. 2. 2. 2.]]

  [[2. 2. 2. 2.]
   [2. 2. 2. 2.]
   [2. 2. 2. 2.]]

  [[2. 2. 2. 2.]
   [2. 2. 2. 2.]
   [2. 2. 2. 2.]]]]

此例中,当调用到Python接口ops.add(x, y)时,会将Python的接口调用通过Pybind11调用到框架的C++层,转换成C++的调用,接着框架会根据用户设置的device_target选择对应的硬件设备,在该硬件设备上执行add这个操作。

从上述原理可以看到,在PyNative模式下,Python脚本代码会根据Python的语法进行执行,而执行过程中涉及到MindSpore的API,会根据用户设置在不同的硬件上进行执行,从而进行加速。因此,在PyNative模式下,用户可以随意使用Python的语法以及调试方法。例如可以使用常见的PyCharm、VS Code等IDE进行代码的调试。

动静结合

JIT

概述

因动态语言的特性以及其灵活高效的开发能力,Python成为AI领域编程的主流语言。但由于Python的解释执行特性,其执行性能往往并非最优。动态图模式贴合Python的解释执行特性,但难以利用算子融合优化等手段进一步优化执行性能。因此MindSpore提供JIT(just-in-time)技术进一步进行性能优化。JIT模式会通过AST树解析的方式或者Python字节码解析的方式,将代码解析为一张中间表示图(IR,intermediate representation)。IR图作为该代码的唯一表示,编译器通过对该IR图的优化,来达到对代码的优化,提高运行性能。与动态图模式相对应,这种JIT的编译模式被称为静态图模式。

基于JIT技术,MindSpore提供了动静结合的方法来提高用户的网络的运行效率。动静结合,即在整体运行为动态图的情况下,指定某些代码块以静态图的方式运行。按照静态图方式运行的代码块会采取先编译后执行的运行模式,在编译期对代码进行全局优化,来获取执行期的性能收益。用户可以通过@jit装饰器修饰函数,来指定其按照静态图的模式执行。有关@jit装饰器的相关文档请见jit API文档。另外,用户还可以通过jit_config来对执行静态图流程的函数进行配置,详情请见mindspore.JitConfig

MindSpore提供了两种JIT编译方式,分别是PIJit以及PSJit。PSJit模式是通过AST树解析的方式,将用户手工标识需要按照PSJit执行的函数转换成静态图,若遇到无法转换成静态图的部分,则会直接报错。PIJit则是通过对Python字节码的解析,在动态图中尽可能的构建静态图,无法转换为静态图的部分则会按照动态图进行执行,来达到动静结合的目的。后续介绍会详细说明二者原理的不同以及各自的特点。

PSJit

在动态图模式下,用户可以通过@jit(mode="PSJit")装饰器修饰函数来让该函数以静态图模式来执行,我们称这种模式为PSJit。同时由于PSJIT为目前jit加速的默认配置,因此也可以直接使用@jit进行装饰。用户需要手动指定需要以静态图模式运行的函数,来获取更加精准的性能收益。与此同时,因为静态图模式需要将函数先进行编译,因此其函数内部使用的语法以及数据结构需要严格遵守静态图严格模式的静态图语法规范。若该函数内存在无法解析的语法或者数据结构,则会直接编译报错。

PSJit的使用方法

用户可以通过@jit装饰器来指定函数以静态图的方式来执行,例如:

[ ]:
import numpy as np
import mindspore as ms
from mindspore import ops
from mindspore import jit
from mindspore import Tensor

@jit
def tensor_cal(x, y, z):
    return ops.matmul(x, y) + z

x = Tensor(np.ones(shape=[2, 3]), ms.float32)
y = Tensor(np.ones(shape=[3, 4]), ms.float32)
z = Tensor(np.ones(shape=[2, 4]), ms.float32)
ret = tensor_cal(x, y, z)
print(ret)
[[4. 4. 4. 4.]
 [4. 4. 4. 4.]]

上述用例中,tensor_cal函数被@jit装饰器修饰,该函数被调用时就会按照静态图的模式进行执行,以获取该函数执行期的性能收益。

PSJit的优点

  • 使用PSJit模式,用户的编程自主性更强,性能优化更精准,可以根据函数特征以及使用经验将网络的性能调至最优。

  • 使用PSJit模式,若遇到静态图内的错误,可以将@jit装饰器删除,以动态图的方式进行程序定位。在问题解决后,再将函数重新指定为静态图模式运行。

PSJit的限制

  • PSJit修饰的函数,其内部的语法必须严格遵守静态图严格模式来进行编程。若使用了静态图不支持的语法或者数据类型,则会直接报错。

PSJit模式的使用建议

  • 相比于动态图执行,被@jit修饰的函数,在第一次调用时需要先消耗一定的时间进行静态图的编译。在该函数的后续调用时,若原有的编译结果可以复用,则会直接使用原有的编译结果进行执行。因此,使用@jit装饰器修饰会多次执行的函数通常会获得更多的性能收益。

  • 静态图模式的运行效率优势体现在其会将被@jit修饰函数进行全局上的编译优化,函数内含有的操作越多,优化的上限也就越高。因此@jit装饰器修饰的函数最好是内含操作很多的大代码块,而不应将很多细碎的、仅含有少量操作的函数分别打上jit标签。否则,则可能会导致性能没有收益甚至劣化。

  • MindSpore静态图绝大部分计算以及优化都是基于对Tensor计算的优化,因此我们建议被修饰的函数应该是那种用来进行真正的数据计算的函数,而不是一些简单的标量计算或者数据结构的变换。

  • @jit修饰的函数,若其输入存在常量,那么该函数每次输入值的变化都会导致重新编译,关于变量常量的概念请见静态图内的常量与变量。因此,建议被修饰的函数以Tensor或者被Mutable修饰的数据作为输入。避免因多次编译导致的额外性能损耗。

PIJit

除了PSJit,MindSpore提供另外一种静态化加速机制PIJit,用户可以通过@jit(mode="PIJit")装饰器修饰函数来让该函数以PIJit模式来执行。当PIJit识别到不支持进入静态图的语法时,会回退到Python执行而非直接编译报错。该功能同时兼顾性能和易用性,减少编译报错的发生。它基于Python字节码的分析,对Python的执行流进行图捕获,让可以以静态图方式运行的子图以静态图方式运行,并让Python语法不支持的子图以动态图方式运行,同时通过修改调整字节码的方式链接动静态图,达到动静混合执行。在满足易用性的前提下,尽可能地提高性能。

PIJit的运行原理

  1. 基于Python虚拟机_PyInterpreterState_SetEvalFrameFunc捕获Python函数的执行,采用上下文管理的方式捕获执行区域内的所有Python函数执行。

  2. 按照当前的运行时输入参数结合函数字节码进行分析,构造控制流图(CFG)以及数据流图(DFG)。

  3. 模拟进栈出栈操作,跟踪逐个字节码,根据栈输入,推导输出。Python3.7~Python3.10每条字节码都有对应的模拟实现,注意是推导输出的类型尺寸,而不是真正执行得到值,除非常量折叠。

  4. 在模拟执行字节码的过程中,将推导结果和操作翻译成MindIR,最后,通过常量折叠,UD分析(删除无用的输入输出参数)等方式,优化静态图。

  5. 在执行等效的静态图之前,对输入参数和优化过程中产生的看护Guard条件进行比对,根据运行时信息,选择匹配的静态图执行。

  6. 动态管理看护Guard和静态图缓冲的匹配关系,对不常用的静态图缓冲进行回收,通过Symbolic Shape和Dynamic Shape优化静态图缓冲。

PIJit的编译流程如下图所示

PIJit的编译流程

PIJit的使用方式

将jit的mode参数设置为PIJit,即可将修饰函数的运行模式切换为PIJit,例如:

[ ]:
import numpy as np
import mindspore as ms
from mindspore import ops
from mindspore import jit
from mindspore import Tensor

@jit(mode="PIJit")
def tensor_cal(x, y, z):
    return ops.matmul(x, y) + z

x = Tensor(np.ones(shape=[2, 3]), ms.float32)
y = Tensor(np.ones(shape=[3, 4]), ms.float32)
z = Tensor(np.ones(shape=[2, 4]), ms.float32)
ret = tensor_cal(x, y, z)
print(ret)
[[4. 4. 4. 4.]
 [4. 4. 4. 4.]]

PIJit的优点

  • 用户体验好,无需人工介入,用户编写的网络代码总是能够正常运行,静态图不能执行的代码会自动采用动态图运行。

  • PIJit可以通过对字节码的变换,使得更多的语句进入静态图。用户无需感知或修改代码。

PIJit的限制

  • 用户无法明确对某些代码做性能加速,对于裂图较多的场景,性能加速的效果可能会不明显。

Shard

在分布式并行场景下,用户可以在动态图模式下调用MindSpore的基础接口构建并行逻辑,但较为复杂。为此,MindSpore提供shard接口,该接口能对Cell或函数进行分布式静态图编译,进而使用静态图的算子并行切分能力。通过shard接口,用户可以在动态图并行的场景下,指定网络的某个Cell或函数以图模式执行并且完成并行操作,实现分布式并行场景的动静结合。该接口只需设置Cell或函数的输入及输出的切分策略,内部通过策略传播机制搜索得到Cell或函数中所有算子的切分策略。当前,本机制还处于实验阶段,详情参见函数式算子切分