Numpy einsum_path 使用方法

在本文中，我们将介绍如何使用Numpy中的einsum_path函数。

einsum_path是Numpy提供的一个非常有用的函数，用于计算多个张量的乘积。它的主要作用是找到最优的计算路径，并返回计算路径的cost和路径本身。

假设我们有两个张量a和b，它们的形状分别为(3, 4)和(4, 5)，我们可以使用如下的代码计算它们的乘积：

import numpy as np

a = np.random.rand(3, 4)
b = np.random.rand(4, 5)

c = np.dot(a, b)

但是，当我们有多个张量需要相乘时，einsum_path就可以发挥它的作用了。

阅读更多：Numpy 教程

einsum_path的使用方法

einsum_path的使用方法非常简单。我们只需要给它传递一个类似于np.einsum的字符串，并指定计算路径的优化等级即可。

具体来说，我们可以使用如下的代码来计算多个张量的乘积，并输出计算路径的cost和路径本身。

import numpy as np

a = np.random.rand(3, 4)
b = np.random.rand(4, 5)
c = np.random.rand(5, 6)

path = np.einsum_path('ij,jk,kl->il', a, b, c, optimize='optimal')

print('cost: {}'.format(path[0]))
print('path: {}'.format(path[1]))

在上面的代码中，我们使用了三个张量a、b、c，并指定了计算路径为’ij,jk,kl->il’。optimize参数指定了计算路径的优化等级，可选的值有’optimal’和’greedy’，其中’optimal’表示使用最优化算法寻找最佳计算路径，’greedy’表示使用贪心算法。

示例

我们可以通过一个简单的实例来演示einsum_path函数的使用方法。

假设我们有三个张量a、b、c，它们的形状分别为(3, 4)、(4, 5)和(5, 6)，我们需要计算它们的乘积。具体来说，我们想要计算下面这个式子：

d(i, j, k) = sum(a(i, l) * b(l, j) * c(k, m))

其中，i、j、k、l、m都是下标。

根据上面的式子，我们可以构造如下的代码来计算d：

import numpy as np

a = np.random.rand(3, 4)
b = np.random.rand(4, 5)
c = np.random.rand(5, 6)

d = np.einsum('il,lm,mj->ijm', a, b, c)

print(d.shape)

执行上面的代码后，我们可以得到一个形状为(3, 4, 6)的张量d。这个张量的每一个元素都是根据式子计算得到的。

但是，我们并不知道上面的代码究竟是如何计算的。如果我们想要知道计算路径，就可以使用einsum_path函数了。

下面是我们使用einsum_path函数的完整代码：

import numpy as np

a = np.random.rand(3, 4)
b = np.random.rand(4, 5)
c = np.random.rand(5, 6)

path = np.einsum_path('il,lm,mj->ijm', a, b, c, optimize='optimal')

print('cost: {}'.format(path[0]))
print('path: {}'.format(path[1]))

d = np.einsum('il,lm,mj->ijm', a, b, c, optimize=path[1])

print(d.shape)

在上面的代码中，我们使用了einsum_path函数，找到了最优的计算路径，并输出了计算路径的cost和路径本身。然后，我们使用einsum函数，并指定了计算路径为path[1]，即最优的计算路径。最后，我们可以得到同样的结果，即形状为(3, 4, 6)的张量d。