函数式和对象式融合编程范式

查看源文件

编程范式(Programming paradigm)是指编程语言的编程风格或编程方式。通常情况下,AI框架均依赖前端编程接口所使用的编程语言的编程范式进行神经网络的构造和训练。MindSpore作为AI+科学计算融合计算框架,分别面向AI、科学计算场景,提供了面向对象编程和函数式编程的支持。同时为提升框架使用的灵活性和易用性,提出了函数式+面向对象融合编程范式,有效地体现了函数式自动微分机制的优势。

下面分别介绍MindSpore支持的三类编程范式及其简单示例。

面向对象编程

面向对象编程(Object-oriented programming, OOP),是指一种将程序分解为封装数据及相关操作的模块(类)而进行的编程方式,对象为类(class)的实例。面向对象编程将对象作为程序的基本单元,将程序和数据封装其中,以提高软件的重用性、灵活性和扩展性,对象里的程序可以访问及经常修改对象相关联的数据。

在一般的编程场景中,代码(code)和数据(data)是两个核心构成部分。面向对象编程是针对特定对象(Object)来设计数据结构,定义类(Class)。类通常由以下两部分构成,分别对应了code和data:

  • 方法(Methods)

  • 属性(Attributes)

对于同一个Class实例化(instantiation)后得到的不同对象而言,方法和属性相同,不同的是属性的值。不同的属性值决定了对象的内部状态,因此OOP能够很好地进行状态管理。

下面为Python构造简单类的示例:

class Sample: #class declaration
    def __init__(self, name): # class constructor (code)
        self.name = name # attribute (data)

    def set_name(self, name): # method declaration (code)
        self.name = name # method implementation (code)

对于构造神经网络来说,首要的组件就是网络层(Layer),一个神经网络层包含以下部分:

  • Tensor操作 (Operation)

  • 权重 (Weights)

此二者恰好与类的Methods和Attributes一一对应,同时权重本身就是神经网络层的内部状态,因此使用类来构造Layer天然符合其定义。此外,我们在编程时希望使用神经网络层进行堆叠,构造深度神经网络,使用OOP编程可以很容易地通过Layer对象组合构造新的Layer类。

下面为使用MindSpore构造神经网络类的示例:

from mindspore import nn, Parameter
from mindspore.common.initializer import initializer

class Linear(nn.Cell):
    def __init__(self, in_features, out_features, has_bias): # class constructor (code)
        super().__init__()
        self.weight = Parameter(initializer('normal', [out_features, in_features], mindspore.float32), 'weight') # layer weight (data)
        self.bias = Parameter(initializer('zeros', [out_features], mindspore.float32), 'bias') # layer weight (data)

    def construct(self, inputs): # method declaration (code)
        output = ops.matmul(inputs, self.weight.transpose(0, 1)) # tensor transformation (code)
        output = output + self.bias # tensor transformation (code)
        return output

除神经网络层的构造使用面向对象编程范式外,MindSpore支持纯面向对象编程方式构造神经网络训练逻辑,此时神经网络的正向计算、反向传播、梯度优化等操作均使用类进行构造。下面是纯面向对象编程的示例:

import mindspore
import mindspore.nn as nn
from mindspore import value_and_grad

class TrainOneStepCell(nn.Cell):
    def __init__(self, network, optimizer):
        super().__init__()
        self.network = network
        self.optimizer = optimizer
        self.grad_fn = value_and_grad(self.network, None, self.optimizer.parameters)

    def construct(self, *inputs):
        loss, grads = self.grad_fn(*inputs)
        self.optimizer(grads)
        return loss

network = nn.Dense(5, 3)
loss_fn = nn.BCEWithLogitsLoss()
network_with_loss = nn.WithLossCell(network, loss_fn)
optimizer = nn.SGD(network.trainable_params(), 0.001)
trainer = TrainOneStepCell(network_with_loss, optimizer)

此时,不论是神经网络,还是其训练过程,均使用继承nn.Cell的类进行管理,可以方便地作为计算图进行编译加速。

函数式编程

函数式编程(Functional programming)是一种将计算机运算视为函数运算,并且避免使用程序状态以及可变对象的编程范式。

在函数式编程中,函数被视为一等公民,这意味着它们可以绑定到名称(包括本地标识符),作为参数传递,并从其他函数返回,就像任何其他数据类型一样。这允许以声明性和可组合的风格编写程序,其中小功能以模块化方式组合。函数式编程有时被视为纯函数式编程的同义词,是将所有函数视为确定性数学函数或纯函数的函数式编程的一个子集。当使用一些给定参数调用纯函数时,它将始终返回相同的结果,并且不受任何可变状态或其他副作用的影响。

函数式编程有两个核心特点,使其十分符合科学计算的需要:

  1. 编程函数语义与数学函数语义完全对等。

  2. 确定性,给定相同输入必然返回相同输出。无副作用。

由于确定性这一特点,通过限制副作用,程序可以有更少的错误,更容易调试和测试,更适合形式验证。

MindSpore提供纯函数式编程的支持,配合mindspore.numpymindspore.scipy提供的数值计算接口,可以便捷地进行科学计算编程。下面是使用函数式编程的示例:

import mindspore.numpy as mnp
from mindspore import grad

grad_tanh = grad(mnp.tanh)
print(grad_tanh(2.0))
# 0.070650816

print(grad(grad(mnp.tanh))(2.0))
print(grad(grad(grad(mnp.tanh)))(2.0))
# -0.13621868
# 0.25265405

配合函数式编程范式的需要,MindSpore提供了多种函数变换接口,涵盖包括自动微分、自动向量化、自动并行、即时编译、数据下沉等功能模块,下面简单进行介绍:

  • 自动微分:gradvalue_and_grad,提供微分函数变换功能;

  • 自动向量化:vmap,用于沿参数轴映射函数 fn 的高阶函数;

  • 自动并行:shard,函数式算子切分,指定函数输入/输出Tensor的分布策略;

  • 即时编译:jit,将Python函数编译为一张可调用的MindSpore图;

  • 数据下沉:data_sink,对输入的函数进行变换,获得可使用数据下沉模式的函数。

基于上述函数变换接口,在使用函数式编程范式时可以快速高效地使用函数变换实现复杂的功能。

函数式微分编程

自动微分

深度学习等现代AI算法,通过使用大量的数据来学习拟合出一个优化后带参模型,其中使用的学习算法,多是基于现实数据自模型中的经验误差来反向传播以更新模型的参数,自动微分技术(Automatic Differentiation,AD)正是其中的关键技术。

自动微分是一种介于数值微分与符号微分之间的一种求导方法。自动微分的核心思想是将计算机程序中的运算操作分解为一个有限的基本操作合集,且合集中基本操作的求导规则均为已知的。在完成每一个基本操作的求导后,使用链式求导法则将结果组合得到整体程序的求导结果。

链式求导法则:

\[(f\circ g)^{'}(x)=f^{'}(g(x))g^{'}(x) \tag{1}\]

根据对分解后的基本操作求导和链式规则的组合不同,自动微分可以分为前向模式反向模式

  • 前向自动微分(Forward Automatic Differentiation,也叫做 tangent linear mode AD)或者前向累加梯度(前向模式)。

  • 后向自动微分(Reverse Automatic Differentiation,也叫做 adjoint mode AD)或者说反向累计梯度(反向模式)。

我们以公式(2)为例介绍前向微分与反向微分的具体计算方式:

\[y=f(x_{1},x_{2})=ln(x_{1})+x_{1}x_{2}-sin(x_{2}) \tag{2}\]

当我们使用前向自动微分求公式(2)在\(x_{1}=2,x_{2}=5\)处的导数 \(\frac{\partial y}{\partial x_{1}}\) 时,前向自动微分的求导方向与原函数的求值方向一致,原函数结果与微分结果可以被同时获得。

forward

当使用反向自动微分时,反向自动微分的求导方向与原函数的求值方向相反,微分结果需要依赖原函数的运行结果。

backward

MindSpore先构建的是基于反向模式的自动微分,并在该方法的基础上实现了正向微分。

为了进一步说明前向微分与反向微分的区别,我们将被求导的原函数,泛化为具有N输入与M输出的函数F:

\[(Y_{1},Y_{2},...,Y_{M})=F(X_{1},X_{2},...,X_{N}) \tag{3}\]

函数 \(F()\) 的导数本身为一个雅可比矩阵(Jacobian matrix)。

\[\begin{split}\left[ \begin{matrix} \frac{\partial Y_{1}}{\partial X_{1}}& ... & \frac{\partial Y_{1}}{\partial X_{N}} \\ ... & ... & ... \\ \frac{\partial Y_{M}}{\partial X_{1}} & ... & \frac{\partial Y_{M}}{\partial X_{N}} \end{matrix} \right] \end{split}\tag{4}\]

前向自动微分

在前向自动微分当中,我们是从输入开始向输出的方向计算的,因此每一次计算我们可以求得输出对某一输入的导数,即雅可比矩阵中的一列。

\[\begin{split}\left[ \begin{matrix} \frac{\partial Y_{1}}{\partial X_{1}}\\ ... \\ \frac{\partial Y_{M}}{\partial X_{1}} \end{matrix} \right] \end{split}\tag{5}\]

为了求取该列的值,自动微分将程序分解为一系列求导规则已知的基本操作,这些基本操作也可以被泛化表达为具有\(n\)输入和\(m\)输出的函数\(f\)

\[(y_{1},y_{2},...,y_{m})=f(x_{1},x_{2},...,x_{n}) \tag{6}\]

由于我们的已知基础函数 \(f\) 的求导规则,即 \(f\) 的雅可比矩阵是已知的。于是我们可以对\(f\)计算雅可比向量积(Jvp, Jacobian-vector-product),并应用链式求导法则获得导数结果。

\[\begin{split}\left[ \begin{matrix} \frac{\partial y_{1}}{\partial X_{i}}\\ ... \\ \frac{\partial y_{m}}{\partial X_{i}} \end{matrix} \right]=\left[ \begin{matrix} \frac{\partial y_{1}}{\partial x_{1}}& ... & \frac{\partial y_{1}}{\partial x_{n}} \\ ... & ... & ... \\ \frac{\partial y_{m}}{\partial x_{1}} & ... & \frac{\partial y_{M}}{\partial x_{n}} \end{matrix} \right]\left[ \begin{matrix} \frac{\partial x_{1}}{\partial X_{i}}\\ ... \\ \frac{\partial x_{n}}{\partial X_{i}} \end{matrix} \right] \end{split}\tag{7}\]

反向自动微分

在反向自动微分当中,我们是从输出开始向输入的方向计算的,因此每一次计算我们可以求得某一输出对输入的导数,即雅可比矩阵中的一行。

\[\begin{split}\left[ \begin{matrix} \frac{\partial Y_{1}}{\partial X_{1}}& ... & \frac{\partial Y_{1}}{\partial X_{N}} \\ \end{matrix} \right] \end{split}\tag{8}\]

为了求取该列的值,自动微分将程序分解为一系列求导规则已知的基本操作,这些基本操作也可以被泛化表达为具有n输入和m输出的函数\(f\)

\[(y_{1},y_{2},...,y_{m})=f(x_{1},x_{2},...,x_{n}) \tag{9}\]

由于我们的已知基础函数\(f\)的求导规则,即f的雅可比矩阵是已知的。于是我们可以对\(f\)计算向量雅可比积(Vjp, Vector-jacobian-product),并应用链式求导法则获得导数结果。

\[\begin{split}\left[ \begin{matrix} \frac{\partial Y_{j}}{\partial x_{1}}& ... & \frac{\partial Y_{j}}{\partial x_{N}} \\ \end{matrix} \right]=\left[ \begin{matrix} \frac{\partial Y_{j}}{\partial y_{1}}& ... & \frac{\partial Y_{j}}{\partial y_{m}} \\ \end{matrix} \right]\left[ \begin{matrix} \frac{\partial y_{1}}{\partial x_{1}}& ... & \frac{\partial y_{1}}{\partial x_{n}} \\ ... & ... & ... \\ \frac{\partial y_{m}}{\partial x_{1}} & ... & \frac{\partial y_{m}}{\partial x_{n}} \end{matrix} \right] \end{split}\tag{10}\]

grad实现

MindSpore中grad使用的是反向自动微分模式,即从正向网络的输出开始计算梯度。

grad算法设计

设模型定义的原函数为:

\[f(g(x, y, z)) \tag{11}\]

\(f()\)\(x\)的梯度为:

\[\frac{df}{dx}=\frac{df}{dg}\frac{dg}{dx}\frac{dx}{dx}+\frac{df}{dg}\frac{dg}{dy}\frac{dy}{dx}+\frac{df}{dg}\frac{dg}{dz}\frac{dz}{dx}\tag{12}\]

\(\frac{df}{dy}\)\(\frac{df}{dz}\)\(\frac{df}{dx}\)类似。

应用链式求导法则,对每个函数(包括算子和图)定义梯度函数bprop: dout->(df, dinputs),这里df表示函数对自由变量(函数外定义的变量)的梯度,dinputs是对函数输入的梯度。在此基础上,应用全微分法则,将(df, dinputs)累加到对应的变量。

MindIR实现了分支,循环,闭包的函数表达式,所以对相应的算子实现正确的反向规则即可求得输入函数的梯度函数。

定义运算符K,反向自动微分算法可以简单表示如下:

v = (func, inputs)
F(v): {
    (result, bprop) = K(func)(K(inputs))
    df, dinputs = bprop(dout)
    v.df += df
    v.dinputs += dinputs
}

grad算法实现

在自动微分流程中,需要进行自动微分的函数会被取出。并作为自动微分模块的输入,并输出对应的梯度图。

MindSpore的自动微分模块实现了从原函数对象到梯度函数对象的转换。转换后的对象为fprop形式的梯度函数对象。

fprop = (forward_result, bprop)forward_result是前向计算图的输出节点, bprop是以fprop的闭包对象形式生成的梯度函数,它只有dout一个入参, inputsoutputs是引用的fprop的输入和输出。

MapObject();    // 实现ValueNode/Parameter/FuncGraph/Primitive对象的映射
MapMorphism();  // 实现CNode的态射
res = k_graph(); // res就是梯度函数的fprop对象

在生成梯度函数对象的过程中,需要完成从原函数到梯度函数的一系列的映射,即为每个原函数中的节点生成其所对应的梯度函数的节点,再按照反向自动微分的规则将这些节点连接在一起,生成梯度函数图。

每张原函数对象的子图都会都会生成一个Dfunctor对象,负责将该原函数对象映射为梯度函数对象。DFunctor主要需要经过 MapObjectMapMorphism两步来实现这种映射关系。

MapObject实现了原函数节点到梯度函数节点的映射,具体包括对自由变量,参数节点以及ValueNode的映射。

MapFvObject();    // 自由变量的映射
MapParamObject(); // 参数节点的映射
MapValueObject(); // ValueNode的映射
  • MapFvObject是对自由变量的映射;

  • MapParamObject是对参数节点的映射;

  • MapValueObject中主要对Primitive以及FuncGraph对象进行映射。

其中,对FuncGraph进行的映射同样需要为该子图创造相应的DFunctor,是一个递归的过程。 Primitive表明了算子的种类,为了支持自动微分,需要为每一种Primitive定义其对应的反向微分函数。

MindSpore将这些定义放在了Python侧,以sin算子为例:

[1]:
import mindspore.ops as ops
from mindspore.ops._grad.grad_base import bprop_getters

@bprop_getters.register(ops.Sin)
def get_bprop_sin(self):
    """Grad definition for `Sin` operation."""
    cos = ops.Cos()

    def bprop(x, out, dout):
        dx = dout * cos(x)
        return (dx,)

    return bprop

x为原函数对象sin的输入,out为原函数对象sin的输出,dout为当前累加的梯度输入。

MapObject完成对以上节点的映射后,MapMorphism从原函数的输出节点开始以递归的方式实现对CNode的态射,建立起节点间的反向传播链接,实现梯度累加。

grad示例

我们构建一个简单的网络来表示公式:

\[f(x) = cos(sin(x)) \tag{13}\]

并对公式(13)的输入x进行求导:

\[f'(x) = -sin(sin(x)) * cos(x) \tag{14}\]

在MindSpore中公式(13)的网络的结构实现为:

[2]:
import mindspore.nn as nn

class Net(nn.Cell):
    def __init__(self):
        super(Net, self).__init__()
        self.sin = ops.Sin()
        self.cos = ops.Cos()

    def construct(self, x):
        a = self.sin(x)
        out = self.cos(a)
        return out

正向网络的结构为:

auto-gradient-foward

对该网络进行反向微分后,所得微分网络结构为:

auto-gradient-forward2

前向自动微分实现

除了支持反向自动微分的grad之外,MindSpore还扩展实现了前向自动微分jvp(Jacobian-Vector-Product)。

相比于反向自动微分,前向自动微分更适合于求取输入维度小于输出维度的网络的梯度。MindSpore的前向自动微分是基于反向自动微分接口grad开发的。

auto-gradient-jvp

黑色为网络的正向流程,第一次求导为针对\(x\)的求导,得到的是蓝色的图。第二次的为蓝色图针对\(v\)的求导,得到的是黄色的图。

黄色的图就是我们所需要的前向模式自动微分的结果图。由于蓝色图可以视为关于\(v\)的线性函数,蓝色节点与黄色节点之间不会存在连边。蓝色节点全部为悬空节点,会被消除,真正运行的就只有原函数节点以及前向微分的节点。因此,该方法不会有额外的运行开销。

参考文献

[1] Baydin, A.G. et al., 2018. Automatic differentiation in machine learning: A survey. arXiv.org. [Accessed September 1, 2021].

函数式+面向对象融合编程

考虑到神经网络模型构建和训练流程的灵活性和易用性需求,结合MindSpore自身的函数式自动微分机制,MindSpore针对AI模型训练设计了函数式+面向对象融合编程范式,可以兼顾面向对象编程和函数式编程的优势,同时使用同一套自动微分机制实现深度学习反向传播和科学计算自动微分的兼容,从底层支持AI和科学计算建模的兼容。下面是函数式+面向对象融合编程的典型过程:

  1. 用类构建神经网络;

  2. 实例化神经网络对象;

  3. 构造正向函数,连接神经网络和损失函数;

  4. 使用函数变换,获得梯度计算(反向传播)函数;

  5. 构造训练过程函数;

  6. 调用函数进行训练。

下面是函数式+面向对象融合编程的简单示例:

# Class definition
class Net(nn.Cell):
    def __init__(self):
        ......
    def construct(self, inputs):
        ......

# Object instantiation
net = Net() # network
loss_fn = nn.CrossEntropyLoss() # loss function
optimizer = nn.Adam(net.trainable_params(), lr) # optimizer

# define forward function
def forword_fn(inputs, targets):
    logits = net(inputs)
    loss = loss_fn(logits, targets)
    return loss, logits

# get grad function
grad_fn = value_and_grad(forward_fn, None, optim.parameters, has_aux=True)

# define train step function
def train_step(inputs, targets):
    (loss, logits), grads = grad_fn(inputs, targets) # get values and gradients
    optimizer(grads) # update gradient
    return loss, logits

for i in range(epochs):
    for inputs, targets in dataset():
        loss = train_step(inputs, targets)

如上述示例,在神经网络构造时,使用面向对象编程,神经网络层的构造方式符合AI编程的习惯。在进行前向计算和反向传播时,MindSpore使用函数式编程,将前向计算构造为函数,然后通过函数变换,获得grad_fn,最后通过执行grad_fn获得权重对应的梯度。

通过函数式+面向对象融合编程,即保证了神经网络构建的易用性,同时提高了前向计算和反向传播等训练过程的灵活性,是MindSpore推荐的默认编程范式。