在线运行下载Notebook下载样例代码查看源文件

计算图

comp-graph.png

计算图是使用有向图表示数学表达式的一种形式。如图所示,神经网络结构可以被视为由Tensor数据和Tensor运算作为节点构成的计算图,因此使用深度学习框架构造神经网络并训练,即是构造计算图,执行计算图的过程。当前在业界框架对计算图的支持分为动态图和静态图两种模式,动态图通过解释执行,具有动态语法亲和性,表达灵活;静态图使用JIT(just in time)编译优化执行,使用静态语法,在语法上有较多限制。

MindSpore同时支持两种计算图模式,采用统一的API表达,在两种模式下使用相同的API,并采用统一的自动微分机制,实现动态图与静态图的统一。下面我们分别介绍MindSpore的两种计算图模式。

动态图

动态图的特点是计算图的构建和计算同时发生(Define by run),其符合Python的解释执行方式,在计算图中定义一个Tensor时,其值就已经被计算且确定,因此在调试模型时较为方便,能够实时得到中间结果的值,但由于所有节点都需要被保存,导致难以对整个计算图进行优化。

在MindSpore中,动态图模式又被称为PyNative模式。由于动态图的解释执行特性,在脚本开发和网络流程调试过程中,推荐使用动态图模式进行调试。

MindSpore默认计算图模式为PyNative模式。PyNative模式下的性能还在持续改善,建议在功能调试中使用,正式生产训练请使用静态图模式。

如需要手动控制框架采用PyNative模式,可以通过以下代码进行配置:

[1]:
import mindspore as ms

ms.set_context(mode=ms.PYNATIVE_MODE)

在PyNative模式下,所有计算节点对应的底层算子均采用单Kernel执行的方式,因此可以任意进行计算结果的打印和调试,如:

[2]:
import numpy as np
from mindspore import nn
from mindspore import ops
from mindspore import Tensor, Parameter

class Network(nn.Cell):
    def __init__(self):
        super().__init__()
        self.w = Parameter(Tensor(np.random.randn(5, 3), ms.float32), name='w') # weight
        self.b = Parameter(Tensor(np.random.randn(3,), ms.float32), name='b') # bias

    def construct(self, x):
        out = ops.matmul(x, self.w)
        print('matmul: ', out)
        out = out + self.b
        print('add bias: ', out)
        return out

model = Network()

我们简单定义一个shape为(5,)的Tensor作为输入,观察输出情况。可以看到在construct方法中插入的print语句将中间结果进行实时的打印输出。

[3]:
x = ops.ones(5, ms.float32)  # input tensor
[4]:
out = model(x)
print('out: ', out)
matmul:  [-1.8809001   2.0400267   0.32370526]
add bias:  [-1.6770952   1.5087128   0.15726662]
out:  [-1.6770952   1.5087128   0.15726662]

静态图

相较于动态图而言,静态图的特点是将计算图的构建和实际计算分开(Define and run)。在构建阶段,根据完整的计算流程对原始的计算图进行优化和调整,编译得到更省内存和计算量更少的计算图。由于编译之后图的结构不再改变,所以称之为 “静态图” 。在计算阶段,根据输入数据执行编译好的计算图得到计算结果。相较于动态图,静态图对全局的信息掌握更丰富,可做的优化也会更多,但是其中间过程对于用户来说是黑盒,无法像动态图一样实时拿到中间计算结果。

在MindSpore中,静态图模式又被称为Graph模式,在Graph模式下,基于图优化、计算图整图下沉等技术,编译器可以针对图进行全局的优化,获得较好的性能,因此比较适合网络固定且需要高性能的场景。

如需要手动控制框架采用Graph模式,可以通过以下代码进行配置:

[5]:
ms.set_context(mode=ms.GRAPH_MODE)

基于源码转换的图编译

在静态图模式下,MindSpore通过源码转换的方式,将Python的源码转换成中间表达IR(Intermediate Representation),并在此基础上对IR图进行优化,最终在硬件设备上执行优化后的图。MindSpore使用基于图表示的函数式IR,称为MindIR,详情可参考中间表示MindIR

MindSpore的静态图执行过程实际包含两步,对应静态图的Define和Run阶段,但在实际使用中,在实例化的Cell对象被调用时并不会感知,MindSpore将两阶段均封装在Cell的__call__方法中,因此实际调用过程为:

model(inputs) = model.compile(inputs) + model.construct(inputs),其中model为实例化Cell对象。

下面我们显式调用compile方法进行示例:

[6]:
model = Network()

model.compile(x)
out = model(x)
print('out: ', out)
...
matmul:
Tensor(shape=[3], dtype=Float32, value=[-4.01971531e+00 -5.79053342e-01  3.41115999e+00])
add bias:
Tensor(shape=[3], dtype=Float32, value=[-3.94732714e+00 -1.46257186e+00  4.50144434e+00])
out:  [-3.9473271 -1.4625719  4.5014443]

静态图语法

在Graph模式下,Python代码并不是由Python解释器去执行,而是将代码编译成静态计算图,然后执行静态计算图。因此,编译器无法支持全量的Python语法。MindSpore的静态图编译器维护了Python常用语法子集,以支持神经网络的构建及训练。详情可参考静态图语法支持

静态图控制流

在PyNative模式下,MindSpore完全支持Python原生语法的流程控制语句。Graph模式下,MindSpore在编译时做了性能优化,因此,在定义网络时使用流程控制语句时会有部分特殊约束,其他部分仍和Python原生语法保持一致。详情可参考流程控制语句

即时编译

通常情况下,由于动态图的灵活性,我们会选择使用PyNative模式来进行自由的神经网络构建,以实现模型的创新和优化。但是当需要进行性能加速时,我们需要对神经网络部分或整体进行加速。此时,直接切换为Graph模式是一种简单选择,但是由于静态图对语法和控制流等限制,使得我们无法从动态图无感知地切换至静态图。

为此,MindSpore提供了jit装饰器,可以通过修饰Python函数或者Python类的成员函数使其被编译成计算图,通过图优化等技术提高运行速度。此时我们可以简单的对想要进行性能优化的模块进行图编译加速,而模型其他部分,仍旧使用解释执行方式,不丢失动态图的灵活性。

Cell模块编译

当我们需要对神经网络的某部分进行加速时,可以直接在construct方法上使用jit修饰器,在调用实例化对象时,该模块自动被编译为静态图。示例如下:

[7]:
import mindspore as ms
from mindspore import nn

class Network(nn.Cell):
    def __init__(self):
        super().__init__()
        self.fc = nn.Dense(10, 1)

    @ms.jit
    def construct(self, x):
        return self.fc(x)

函数编译

与Cell模块编译类似,在需要对Tensor的某些运算进行编译加速时,可以在其定义的函数上使用jit修饰器,在调用该函数时,该模块自动被编译为静态图。示例如下:

基于MindSpore的函数式自动微分特性,推荐使用函数编译方式对Tensor运算进行即时编译加速。

[8]:
@ms.jit
def mul(x, y):
    return x * y

整图编译

MindSpore支持将神经网络训练的正向计算、反向传播、梯度优化更新等步骤合为一个计算图进行编译优化,此方法称为整图编译。此时,仅需将神经网络训练逻辑构造为函数,并在函数上使用jit修饰器,即可达到整图编译的效果。

下面使用简单的全连接网络进行举例:

[9]:
network = nn.Dense(10, 1)
loss_fn = nn.BCELoss()
optimizer = nn.Adam(network.trainable_params(), 0.01)

def forward_fn(data, label):
    logits = network(data)
    loss = loss_fn(logits, label)
    return loss

grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters)

@ms.jit
def train_step(data, label):
    loss, grads = grad_fn(data, label)
    optimizer(grads)
    return loss

如上述代码所示,将神经网络正向执行与损失函数封装为forward_fn后,执行函数变换获得梯度计算函数。而后将梯度计算函数、优化器调用封装为train_step函数,并使用jit进行修饰,调用train_step函数时,会进行静态图编译,获得整图并执行。

除使用修饰器外,也可使用函数变换方式调用jit方法,示例如下:

[10]:
train_step = ms.jit(train_step)