MindSpore尝鲜之爱因斯坦求和

技术背景

在前面的博客中,我们介绍过关于numpy中的张量网络的一些应用,同时利用相关的张量网络操作,我们可以实现一些分子动力学模拟中的约束算法,如LINCS等。在最新的nightly版本的MindSpore中也支持了爱因斯坦求和的算子,这是在张量网络中非常核心的一个操作,本文就简单介绍一下MindSpore中使用爱因斯坦求和的方法。

安装最新版的MindSpore

Einsum是在1.6之后的版本才支持的,MindSpore的Master分支就是官网上面的Nightly版本,我们可以安装这个已经实现了爱因斯坦求和算子的版本。

安装指令如下:

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

简单示例

我们可以先将张量缩并,理解成是一个普通的矩阵乘积即可,只是相乘的矩阵维度大小略有区别。可以先看一个简单的案例,再解析其中的原理:

Python 3.9.0 (default, Nov 15 2020, 14:28:56)  [GCC 7.3.0] :: Anaconda, Inc. on linux Type "help", "copyright", "credits" or "license" for more information. >>> import os >>> os.environ['GLOG_v'] = '4' # 设定日志等级,防止屏幕上出现大量的MindSpore告警信息 >>> from mindspore import Tensor, ops >>> a = Tensor([[1.,2.],[3.,4.]]) >>> b = Tensor([1.,1.]) >>> ein_0 = ops.Einsum('ij,j->i') # 对第二个维度进行缩并 >>> print (ein_0((a,b))) [3. 7.] >>> ein_1 = ops.Einsum('ij,i->j') # 对第一个维度进行缩并 >>> print (ein_1((a,b))) [4. 6.] 

原理解析

我们日常所见的矩阵,可以采用张量这样的“章鱼图”表示方法来标记,每一个张量都是一个“章鱼”的头,而矩阵的每一个维度代表一条“章鱼腿”,比如一个维度为/((2,2,2)/)的矩阵,就可以用一只“三条腿的章鱼”来表示。而学习张量网络的时候经常可以看到的如下这张图,就分别用于表示一维、二维和三维的矩阵:

这些是张量的基本概念,而如果我们把“章鱼腿”都连接起来,就表示规定了这个张量的运算方向,张量只能跟相互连接的张量进行缩并(矩阵运算)操作。比如上一个章节中的案例,就可以用这样的一个张量图来表示:

可以看到,这两个矩阵运算之后得到的结果,是一个“单腿”的张量,对应于矩阵运算的维度变化就是:/((2,2)/times(2,)->(2,)/)。如果我们把这个运算推广开来,建立一个带有复杂链接信息的张量拓扑图,那么就得到了一个张量网络。在一个大型的张量网络中,我们不仅可以通过高性能的GPU来加速张量缩并的运算,还有路径搜索和图分解等算法可以进一步加速张量网络缩并,得到我们想要的结果。

我们可以再细化的讲解一下上一个章节中的案例,当我们使用ij,j->i这个路径时,得到的结果的第一个元素为/(a[0]/cdot b=3/),得到的第二个元素为/(a[1]/cdot b=7/)。当我们使用ij,i->j这个路径时,得到的结果的第一个元素为/(a[:,0]/cdot b=4/),得到的第二个元素为:/(a[:,1]/cdot b=4/)。从这个过程可以发现,对于二维的张量,其实取不同的“腿”,就是取行和取列运算的区别而已。对于更加高维的张量,我们很难用行、列这样的指标来衡量和理解,但是我们还是可以通过python的这种矩阵运算来理解。更多的参考示例,可以阅读MindSpore的官方文档(参考链接2)。

总结概要

张量网络计算,已经在众多的领域中得到了应用,不仅仅是传统的计算化学,当下医药研发领域的分子动力学模拟、计算化学和材料模拟,甚至是未来的量子计算,张量网络技术都在当中发挥重要作用。本文介绍的是MindSpore最新对张量网络计算的支持的第一步:用爱因斯坦求和计算张量网络缩并。

版权声明

本文首发链接为:https://www.gaodi.net/dechinphy/p/mseinsum.html

作者ID:DechinPhy

标签:

商匡云商
Logo
注册新帐户
对比商品
  • 合计 (0)
对比
0
购物车