昇思MindSpore学习笔记10--使用静态图加速

摘要：

昇思MindSpore AI框架支持动态图、静态图两种模式。默认动态图模式便于调试，静态图模式用于加速改善性能场景。开启静态图有jit装饰器和set_context方法2种开启方式。

一、概念

AI编译框架两种运行模式

动态图模式(PyNative模式，默认选项)

计算图构建和计算同时发生（Define by run）

计算图在定义Tensor时，其值已经确定

便于Python调试模型，可实时得到中间结果。

由于要保存所有节点，优化困难。

静态图模式(Graph模式)

计算图的构建和实际计算分开（Define and run）

基于图优化、计算图整图下沉等技术，对图进行全局优化，性能较好

适合神经网络固定、高性能场景。

张量Tensor数据的计算以及其微分处理

反复执行

性能加速

部分Python语法不支持。参考静态图语法支持

二、环境准备

安装minspore模块

!pip uninstall mindspore -y
!pip install -i https://pypi.mirrors.ustc.edu.cn/simple mindspore==2.3.0rc1

导入numpy、minspore、nn、Tensor等相关模块

import numpy as np
import mindspore as ms
from mindspore import nn, Tensor

三、动态图模式(PyNative模式)

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 logitsmodel = Network()
input = Tensor(np.ones([64, 1, 28, 28]).astype(np.float32))
output = model(input)
print(output)

输出：

[[-0.00134926 -0.13563682 -0.02863023 -0.05452826  0.03290743 -0.12423715-0.0582641  -0.10854103 -0.08558805  0.06099342][-0.00134926 -0.13563682 -0.02863023 -0.05452826  0.03290743 -0.12423715-0.0582641  -0.10854103 -0.08558805  0.06099342][-0.00134926 -0.13563682 -0.02863023 -0.05452826  0.03290743 -0.12423715-0.0582641  -0.10854103 -0.08558805  0.06099342][-0.00134926 -0.13563682 -0.02863023 -0.05452826  0.03290743 -0.12423715-0.0582641  -0.10854103 -0.08558805  0.06099342][-0.00134926 -0.13563682 -0.02863023 -0.05452826  0.03290743 -0.12423715-0.0582641  -0.10854103 -0.08558805  0.06099342]...[-0.00134926 -0.13563682 -0.02863023 -0.05452826  0.03290743 -0.12423715-0.0582641  -0.10854103 -0.08558805  0.06099342][-0.00134926 -0.13563682 -0.02863023 -0.05452826  0.03290743 -0.12423715-0.0582641  -0.10854103 -0.08558805  0.06099342][-0.00134926 -0.13563682 -0.02863023 -0.05452826  0.03290743 -0.12423715-0.0582641  -0.10854103 -0.08558805  0.06099342][-0.00134926 -0.13563682 -0.02863023 -0.05452826  0.03290743 -0.12423715-0.0582641  -0.10854103 -0.08558805  0.06099342]]

四、静态图模式

下例手动控制框架采用静态图模式。

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 logitsmodel = Network()
input = Tensor(np.ones([64, 1, 28, 28]).astype(np.float32))
output = model(input)
print(output)

输出：

[[ 0.05363735  0.05117104 -0.03343301  0.06347139  0.07546629  0.032630910.02790363  0.06269836  0.01838502  0.04387159][ 0.05363735  0.05117104 -0.03343301  0.06347139  0.07546629  0.032630910.02790363  0.06269836  0.01838502  0.04387159][ 0.05363735  0.05117104 -0.03343301  0.06347139  0.07546629  0.032630910.02790363  0.06269836  0.01838502  0.04387159][ 0.05363735  0.05117104 -0.03343301  0.06347139  0.07546629  0.032630910.02790363  0.06269836  0.01838502  0.04387159]...[ 0.05363735  0.05117104 -0.03343301  0.06347139  0.07546629  0.032630910.02790363  0.06269836  0.01838502  0.04387159][ 0.05363735  0.05117104 -0.03343301  0.06347139  0.07546629  0.032630910.02790363  0.06269836  0.01838502  0.04387159][ 0.05363735  0.05117104 -0.03343301  0.06347139  0.07546629  0.032630910.02790363  0.06269836  0.01838502  0.04387159][ 0.05363735  0.05117104 -0.03343301  0.06347139  0.07546629  0.032630910.02790363  0.06269836  0.01838502  0.04387159]]

五、静态图模式开启方式

1.基于jit装饰器的开启方式

示例：

import numpy as np
import mindspore as ms
from mindspore import nn, Tensorclass 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 logitsinput = 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)

输出：

[[-0.12126954  0.06986676 -0.2230821  -0.07087803 -0.01003947  0.010633920.10143848 -0.0200909  -0.09724037  0.0114444 ][-0.12126954  0.06986676 -0.2230821  -0.07087803 -0.01003947  0.010633920.10143848 -0.0200909  -0.09724037  0.0114444 ][-0.12126954  0.06986676 -0.2230821  -0.07087803 -0.01003947  0.010633920.10143848 -0.0200909  -0.09724037  0.0114444 ][-0.12126954  0.06986676 -0.2230821  -0.07087803 -0.01003947  0.010633920.10143848 -0.0200909  -0.09724037  0.0114444 ]...[-0.12126954  0.06986676 -0.2230821  -0.07087803 -0.01003947  0.010633920.10143848 -0.0200909  -0.09724037  0.0114444 ][-0.12126954  0.06986676 -0.2230821  -0.07087803 -0.01003947  0.010633920.10143848 -0.0200909  -0.09724037  0.0114444 ][-0.12126954  0.06986676 -0.2230821  -0.07087803 -0.01003947  0.010633920.10143848 -0.0200909  -0.09724037  0.0114444 ][-0.12126954  0.06986676 -0.2230821  -0.07087803 -0.01003947  0.010633920.10143848 -0.0200909  -0.09724037  0.0114444 ]]

使用修饰器函数变换方式调用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)

输出：

[[ 0.11027216 -0.09628229  0.0457969   0.05396656 -0.06958974  0.0428197-0.1572069  -0.14151613 -0.04531277  0.07521383][ 0.11027216 -0.09628229  0.0457969   0.05396656 -0.06958974  0.0428197-0.1572069  -0.14151613 -0.04531277  0.07521383][ 0.11027216 -0.09628229  0.0457969   0.05396656 -0.06958974  0.0428197-0.1572069  -0.14151613 -0.04531277  0.07521383][ 0.11027216 -0.09628229  0.0457969   0.05396656 -0.06958974  0.0428197-0.1572069  -0.14151613 -0.04531277  0.07521383]...[ 0.11027216 -0.09628229  0.0457969   0.05396656 -0.06958974  0.0428197-0.1572069  -0.14151613 -0.04531277  0.07521383][ 0.11027216 -0.09628229  0.0457969   0.05396656 -0.06958974  0.0428197-0.1572069  -0.14151613 -0.04531277  0.07521383][ 0.11027216 -0.09628229  0.0457969   0.05396656 -0.06958974  0.0428197-0.1572069  -0.14151613 -0.04531277  0.07521383][ 0.11027216 -0.09628229  0.0457969   0.05396656 -0.06958974  0.0428197-0.1572069  -0.14151613 -0.04531277  0.07521383]]

对神经网络的某部分进行加速时，直接在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)

输出：

[[ 0.10522258  0.06597593 -0.09440921 -0.04883489  0.07194916  0.1343117-0.06813788  0.01986085  0.0216996  -0.05345828][ 0.10522258  0.06597593 -0.09440921 -0.04883489  0.07194916  0.1343117-0.06813788  0.01986085  0.0216996  -0.05345828][ 0.10522258  0.06597593 -0.09440921 -0.04883489  0.07194916  0.1343117-0.06813788  0.01986085  0.0216996  -0.05345828][ 0.10522258  0.06597593 -0.09440921 -0.04883489  0.07194916  0.1343117-0.06813788  0.01986085  0.0216996  -0.05345828]...[ 0.10522258  0.06597593 -0.09440921 -0.04883489  0.07194916  0.1343117-0.06813788  0.01986085  0.0216996  -0.05345828][ 0.10522258  0.06597593 -0.09440921 -0.04883489  0.07194916  0.1343117-0.06813788  0.01986085  0.0216996  -0.05345828][ 0.10522258  0.06597593 -0.09440921 -0.04883489  0.07194916  0.1343117-0.06813788  0.01986085  0.0216996  -0.05345828][ 0.10522258  0.06597593 -0.09440921 -0.04883489  0.07194916  0.1343117-0.06813788  0.01986085  0.0216996  -0.05345828]]

2.基于全局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)

输出：

[[ 0.08501796 -0.04404321 -0.05165704  0.00357929  0.00051521  0.009464560.02748473 -0.19415936 -0.00278988  0.04024826][ 0.08501796 -0.04404321 -0.05165704  0.00357929  0.00051521  0.009464560.02748473 -0.19415936 -0.00278988  0.04024826][ 0.08501796 -0.04404321 -0.05165704  0.00357929  0.00051521  0.009464560.02748473 -0.19415936 -0.00278988  0.04024826][ 0.08501796 -0.04404321 -0.05165704  0.00357929  0.00051521  0.009464560.02748473 -0.19415936 -0.00278988  0.04024826]...[ 0.08501796 -0.04404321 -0.05165704  0.00357929  0.00051521  0.009464560.02748473 -0.19415936 -0.00278988  0.04024826][ 0.08501796 -0.04404321 -0.05165704  0.00357929  0.00051521  0.009464560.02748473 -0.19415936 -0.00278988  0.04024826][ 0.08501796 -0.04404321 -0.05165704  0.00357929  0.00051521  0.009464560.02748473 -0.19415936 -0.00278988  0.04024826][ 0.08501796 -0.04404321 -0.05165704  0.00357929  0.00051521  0.009464560.02748473 -0.19415936 -0.00278988  0.04024826]]