MindSpore尝鲜之Vmap功能

技术背景

Vmap是一种在python里面经常提到的向量化运算的功能,比如之前大家常用的就是numba和jax中的向量化运算的接口。虽然numpy中也使用到了向量化的运算,比如计算两个numpy数组的加和,就是一种向量化的运算。但是在numpy中模块封装的较好,定制化程度低,但是使用便捷,只需要调用最上层的接口即可。现在最新版本的mindspore也已经推出了vmap的功能,像mindspore、numba还有jax,与numpy的最大区别就是,需要在使用过程中对需要向量化运算的函数额外嵌套一层vmap的函数,这样就可以实现只对需要向量化运算的模块进行扩展。用一个公式来理解向量化运算的话就是:

\[a_1+b_1=c_1\\ a_2+b_2=c_2\\ .\\ .\\ .\\ a_n+b_n=c_n\\ \Downarrow\\ \vec{a}+\vec{b}=\vec{c} \]

安装最新版MindSpore

关于jax中的vmap使用案例,可以参考前面介绍的LINCS约束算法实现SETTLE约束算法批量化实现这两篇文章,都有使用到jax的vmap功能,这里我们着重介绍的是MindSpore中最新实现的vmap功能。首先我们需要安装mindspore最新的Nightly版本,其对应的是MindSpore的Gitee仓库中的master分支,具体安装指令可以参考其官方链接

因为我们本地已经安装过Mindspore的旧版本,因此还需要在安装指令之后加上--upgrade操作,否则会导致系统误以为本地已经安装成功,不会执行安装的操作:

$ python3 -m pip install mindspore-cuda11-dev -i https://pypi.tuna.tsinghua.edu.cn/simple --upgrade

Vmap功能测试

这里我们先来看一个比较简单的示例:

In [1]: from mindspore import Tensor

In [2]: from mindspore.ops.functional import vmap

In [3]: y = lambda a,b: a+b

In [4]: A = Tensor([1,2,3])

In [5]: B = Tensor([3,4,5])

In [6]: vmap_y = vmap(y,in_axes=(0,0))

In [7]: y(A[0],B[0]) # 元素加和
Out[7]: Tensor(shape=[], dtype=Int64, value= 4)

In [8]: vmap_y(A,B) # 矢量加和
Out[8]: Tensor(shape=[3], dtype=Int64, value= [4, 6, 8])

在上面的这个示例中,我们定义了一个加法函数y,作用就是把输入的两个对象相加。这里需要注意的是,如果输入给y的是两个Mindspore的Tensor对象,那么会直接返回两个Tensor对应位置相加的结果。但是如果输入给y的是两个普通python的list,则输出的结果会是两个list的拼接,这跟不同类型的加法的实现方式有关,在文末总结中会进行解释。这里我们只是想说明:y本身是一个元素加和的函数,可以通过vmap使其称为矢量加和的函数。关于输入的in_axes参数,指的是扩展的维度。比如我们写了一个支持\((A,A)\times(A,1)\)维度的函数,如果把in_axes参数设置为0,那么就可以得到一个支持计算\((B,A,A)\times(B,A,1)\)维度的函数。其中in_axes参数,决定的是被扩展的维度B所在的位置。这一点我们可以看一下vmap的官方示例:

在这个案例中,也是定义了一个普通的加和函数,通过vmap去扩展不同的维度,大致的计算逻辑为:

\[(A,)+(A,)+(A,)\\ \Downarrow^{in\_axes=(0,1,None)}\\ (B,A)+(A,B)+(A,)=(B,A)+(B,A)+(1,A)=(B,A)\\ \Downarrow^{out\_axes=1}\\ (A,B) \]

其实这个过程中关于in_axes是比较容易可以理解的,但是这个out_axes有时候会让人难以捉摸,在github上专门有人提出了这个issue并有人做出了解释:

结合上面的案例,其实out_axes就是决定了扩展的维度B在结果中的位置,比如out_axes=1,所对应的结果中就是\((x,B,x,...x)\)。也就是说,其不影响计算的结果,但是有可能会对计算结果进行转置操作,在MindSpore和Numpy中称为swap_axes

总结概要

本文介绍了华为推出的深度学习框架MindSpore中最新支持的vmap功能函数,可以用于向量化的计算,本质上的主要作用是替代并加速python中的for循环的操作。最早是在numba和pytroch、jax中对vmap功能进行了支持,其实numpy中的底层计算也用到了向量化的运算,因此速度才如此之快。vmap在python中更多的是与即时编译功能jit一同使用,能够起到简化编程的同时对性能进行极大程度的优化,尤其是python中的for循环的优化。但是对于一些numpy、jax或者MindSpore中已有的算子而言,还是建议直接使用其已经实现的算子,而不是vmap再手写一个。

版权声明

本文首发链接为:https://www.cnblogs.com/dechinphy/p/ms-vmap.html

作者ID:DechinPhy

更多原著文章请参考:https://www.cnblogs.com/dechinphy/

打赏专用链接:https://www.cnblogs.com/dechinphy/gallery/image/379634.html

腾讯云专栏同步:https://cloud.tencent.com/developer/column/91958

参考链接

  1. https://gitee.com/mindspore/mindspore/blob/master/mindspore/python/mindspore/ops/functional.py#L845
posted @ 2022-03-28 11:08  DECHIN  阅读(358)  评论(0编辑  收藏  举报