学习内容:使用静态图加速
1.背景介绍
AI编译框架分为两种运行模式,分别是动态图模式以及静态图模式。MindSpore默认情况下是以动态图模式运行,但也支持手工切换为静态图模式。
-
动态图的特点
是计算图的构建和计算同时发生(Define by run),其符合Python的解释执行方式,在计算图中定义一个Tensor时,其值就已经被计算且确定,因此在调试模型时较为方便,能够实时得到中间结果的值,但由于所有节点都需要被保存,导致难以对整个计算图进行优化。 -
静态图的特点
相较于动态图而言,静态图的特点是将计算图的构建和实际计算分开(Define and run),
在MindSpore中,静态图模式又被称为Graph模式,在Graph模式下,基于图优化、计算图整图下沉等技术,编译器可以针对图进行全局的优化,获得较好的性能,因此比较适合网络固定且需要高性能的场景。
2.动态图模式
在MindSpore中,动态图模式又被称为PyNative模式。由于动态图的解释执行特性,在脚本开发和网络流程调试过程中,推荐使用动态图模式进行调试。
import numpy as np
import mindspore as ms
from mindspore import nn, Tensor
ms.set_context(mode=ms.PYNATIVE_MODE) # 使用set_context进行动态图模式的配置
class Network(nn.Cell):
def __init__(self):
super().__init__()
self.flatten = nn.Flatten()
self.dense_relu_sequential = nn.SequentialCell(
nn.Dense(28*28, 512),
nn.ReLU(),
nn.Dense(512, 512),
nn.ReLU(),
nn.Dense(512, 10)
)
def construct(self, x):
x = self.flatten(x)
logits = self.dense_relu_sequential(x)
return logits
model = Network()
input = Tensor(np.ones([64, 1, 28, 28]).astype(np.float32))
output = model(input)
print(output)
3.静态图模式
- 在Graph模式下,Python代码并不是由Python解释器去执行,而是将代码编译成静态计算图,然后执行静态计算图。
MindSpore的静态图执行过程实际包含两步,对应静态图的Define和Run阶段,但在实际使用中,在实例化的Cell对象被调用时用户并不会分别感知到这两阶段,MindSpore将两阶段均封装在Cell的__call__方法中,因此实际调用过程为:
model(inputs) = model.compile(inputs) + model.construct(inputs)
其中model为实例化Cell对象。
需要手动控制框架采用静态图模式,可以通过以下代码进行网络构建:
import numpy as np
import mindspore as ms
from mindspore import nn, Tensor
ms.set_context(mode=ms.GRAPH_MODE) # 使用set_context进行运行静态图模式的配置
class Network(nn.Cell):
def __init__(self):
super().__init__()
self.flatten = nn.Flatten()
self.dense_relu_sequential = nn.SequentialCell(
nn.Dense(28*28, 512),
nn.ReLU(),
nn.Dense(512, 512),
nn.ReLU(),
nn.Dense(512, 10)
)
def construct(self, x):
x = self.flatten(x)
logits = self.dense_relu_sequential(x)
return logits
model = Network()
input = Tensor(np.ones([64, 1, 28, 28]).astype(np.float32))
output = model(input)
print(output)
- 静态图模式的使用场景
MindSpore编译器重点面向Tensor数据的计算以及其微分处理。因此使用MindSpore API以及基于Tensor对象的操作更适合使用静态图编译优化。其他操作虽然可以部分入图编译,但实际优化作用有限。另外,静态图模式先编译后执行的模式导致其存在编译耗时。因此,如果函数无需反复执行,那么使用静态图加速也可能没有价值。 - 静态图模式开启方式
通常情况下,由于动态图的灵活性,我们会选择使用PyNative模式来进行自由的神经网络构建,以实现模型的创新和优化。但是当需要进行性能加速时,我们需要对神经网络部分或整体进行加速。MindSpore提供了两种切换为图模式的方式,分别是基于装饰器的开启方式以及基于全局context的开启方式。 - 基于装饰器的开启方式
MindSpore提供了jit装饰器,可以**通过修饰Python函数或者Python类的成员函数使其被编译成计算图,通过图优化等技术提高运行速度。**此时我们可以简单的对想要进行性能优化的模块进行图编译加速,而模型其他部分,仍旧使用解释执行方式,不丢失动态图的灵活性。无论全局context是设置成静态图模式还是动态图模式,被jit修饰的部分始终会以静态图模式进行运行。
在需要对Tensor的某些运算进行编译加速时,可以在其定义的函数上使用jit修饰器,在调用该函数时,该模块自动被编译为静态图。需要注意的是,jit装饰器只能用来修饰函数,无法对类进行修饰。jit的使用示例如下:
import numpy as np
import mindspore as ms
from mindspore import nn, Tensor
class Network(nn.Cell):
def __init__(self):
super().__init__()
self.flatten = nn.Flatten()
self.dense_relu_sequential = nn.SequentialCell(
nn.Dense(28*28, 512),
nn.ReLU(),
nn.Dense(512, 512),
nn.ReLU(),
nn.Dense(512, 10)
)
def construct(self, x):
x = self.flatten(x)
logits = self.dense_relu_sequential(x)
return logits
input = Tensor(np.ones([64, 1, 28, 28]).astype(np.float32))
@ms.jit # 使用ms.jit装饰器,使被装饰的函数以静态图模式运行
def run(x):
model = Network()
return model(x)
output = run(input)
print(output)
除使用修饰器外,也可使用函数变换方式调用jit方法,示例如下:
import numpy as np
import mindspore as ms
from mindspore import nn, Tensor
class Network(nn.Cell):
def __init__(self):
super().__init__()
self.flatten = nn.Flatten()
self.dense_relu_sequential = nn.SequentialCell(
nn.Dense(28*28, 512),
nn.ReLU(),
nn.Dense(512, 512),
nn.ReLU(),
nn.Dense(512, 10)
)
def construct(self, x):
x = self.flatten(x)
logits = self.dense_relu_sequential(x)
return logits
input = Tensor(np.ones([64, 1, 28, 28]).astype(np.float32))
def run(x):
model = Network()
return model(x)
run_with_jit = ms.jit(run) # 通过调用jit将函数转换为以静态图方式执行
output = run(input)
print(output)
当我们需要对神经网络的某部分进行加速时,可以直接在construct方法上使用jit修饰器,在调用实例化对象时,该模块自动被编译为静态图。示例如下:
import numpy as np
import mindspore as ms
from mindspore import nn, Tensor
class Network(nn.Cell):
def __init__(self):
super().__init__()
self.flatten = nn.Flatten()
self.dense_relu_sequential = nn.SequentialCell(
nn.Dense(28*28, 512),
nn.ReLU(),
nn.Dense(512, 512),
nn.ReLU(),
nn.Dense(512, 10)
)
@ms.jit # 使用ms.jit装饰器,使被装饰的函数以静态图模式运行
def construct(self, x):
x = self.flatten(x)
logits = self.dense_relu_sequential(x)
return logits
input = Tensor(np.ones([64, 1, 28, 28]).astype(np.float32))
model = Network()
output = model(input)
print(output)
- 基于context的开启方式
context模式是一种全局的设置模式。代码示例如下:
import numpy as np
import mindspore as ms
from mindspore import nn, Tensor
ms.set_context(mode=ms.GRAPH_MODE) # 使用set_context进行运行静态图模式的配置
class Network(nn.Cell):
def __init__(self):
super().__init__()
self.flatten = nn.Flatten()
self.dense_relu_sequential = nn.SequentialCell(
nn.Dense(28*28, 512),
nn.ReLU(),
nn.Dense(512, 512),
nn.ReLU(),
nn.Dense(512, 10)
)
def construct(self, x):
x = self.flatten(x)
logits = self.dense_relu_sequential(x)
return logits
model = Network()
input = Tensor(np.ones([64, 1, 28, 28]).astype(np.float32))
output = model(input)
print(output)
4.静态图的语法约束
在Graph模式下,Python代码并不是由Python解释器去执行,而是将代码编译成静态计算图,然后执行静态计算图。因此,编译器无法支持全量的Python语法。MindSpore的静态图编译器维护了Python常用语法子集,以支持神经网络的构建及训练。
基础语法的语法约束
图模式下的执行图是从源码转换而来,并不是所有的Python语法都能支持。
- 当construct函数里,使用未定义的类成员时,将抛出AttributeError异常。示例如下:
import mindspore as ms
from mindspore import nn, set_context
set_context(mode=ms.GRAPH_MODE)
class Net(nn.Cell):
def __init__(self):
super(Net, self).__init__()
def construct(self, x):
return x + self.y
net = Net()
net(1)
结果报错如下:
AttributeError: External object has no attribute y
- nn.Cell不支持classmethod修饰的类方法。示例如下:
import mindspore as ms
ms.set_context(mode=ms.GRAPH_MODE)
class Net(ms.nn.Cell):
@classmethod
def func(cls, x, y):
return x + y
def construct(self, x, y):
return self.func(x, y)
net = Net()
out = net(ms.Tensor(1), ms.Tensor(2))
print(out)
结果报错如下:
TypeError: The parameters number of the function is 3, but the number of provided arguments is 2.
- 在图模式下,有些Python语法难以转换成图模式下的中间表示MindIR。对标Python的关键字,存在部分关键字在图模式下是不支持的:AsyncFunctionDef、Delete、AnnAssign、AsyncFor、AsyncWith、Match、Try、Import、ImportFrom、Nonlocal、NamedExpr、Set、SetComp、Await、Yield、YieldFrom、Starred。如果在图模式下使用相关的语法,将会有相应的报错信息提醒用户。
import mindspore as ms
@ms.jit
def test_try_except(x, y):
global_out = 1
try:
global_out = x / y
except ZeroDivisionError:
print("division by zero, y is zero.")
return global_out
test_try_except_out = test_try_except(1, 0)
print("out:", test_try_except_out)
结果报错如下:
RuntimeError: Unsupported statement 'Try'.
扩展语法(LAX级别)
调用第三方库
- 第三方库
- Python内置模块和Python标准库。例如os、sys、math、time等模块。
- 第三方代码库。路径在Python安装目录的site-packages目录下,需要先安装后导入,例如NumPy、SciPy等。需要注意的是,mindyolo、mindflow等MindSpore套件不被视作第三方库,具体列表可以参考parser文件的 _modules_from_mindspore 列表。
- 通过环境变量MS_JIT_IGNORE_MODULES指定的模块。与之相对的有环境变量MS_JIT_MODULES,具体使用方法请参考环境变量。
- 支持第三方库的数据类型,允许调用和返回第三方库的对象。
- 支持调用第三方库的方法。
- 支持使用NumPy第三方库数据类型创建Tensor对象。
- 暂不支持对第三方库数据类型的下标索引赋值。
- 支持自定义类的使用:支持在图模式下使用用户自定义的类,可以对类进行实例化,使用对象的属性及方法。
5.JitConfig配置选项
在图模式下,可以通过使用JitConfig配置选项来一定程度的自定义编译流程,目前JitConfig支持的配置参数如下: