如何使用numpy.einsum对未知数量的操作数进行矩阵向量乘法?

llew8vvj  于 2023-06-29  发布在  其他
关注(0)|答案(3)|浏览(79)

我想在Python中高效地执行一系列矩阵向量乘法,numpy.einsum函数似乎是最佳选择。然而,我不知道链中矩阵操作数 N 的数量,因此没有简单的方法来写subscripts字符串传递给einsum
为了更清楚,让我们考虑 N=3 矩阵 ABC 以及初始向量v的简单情况。numpy.einsum调用如下:

result = numpy.einsum('ij,jk,kl,l', A, B, C, v)

如何将此操作推广到任意 N?一种可能性是以编程方式创建subscripts字符串,但numpy.einsum只支持52个不同的标签(26个字母,小写和大写),这在我的情况下代表了一个限制,因为我事先不知道操作数的数量。
是否有一种方法可以通过使用基于不同numpy函数的另一种方法来尽可能有效地做到这一点?

nc1teljy

nc1teljy1#

这里有两个递归解决方案,一个基于“原子归纳法”,也就是对每个矩阵乘法做一个einsum:

A, B, C, D = (np.random.rand(10,10),) * 4

def recursive_matrix_multiply(*args):
    if len(args) == 1:
        return args[0]
    return recursive_matrix_multiply( np.einsum('ij,jk' if args[1].ndim == 2 else 'ij,j', *args[:2]), *args[2:])

np.testing.assert_almost_equal( recursive_matrix_multiply(A, B, C, D), np.einsum('ij,jk,kl,lm', A, B, C, D))
np.testing.assert_almost_equal( recursive_matrix_multiply(A, B, C, D[:,0]), np.einsum('ij,jk,kl,l', A, B, C, D[:,0]))

这里是一个基于“块归纳”的解决方案,在1处对每个最大块进行einsum。请注意,我从string.ascii_letters开始使用函数vocab中的所有52个字母,但我遇到了错误,所以我不得不将其减少到1115确实工作,但速度很慢。

from string import ascii_letters, ascii_lowercase
from itertools import pairwise
from functools import reduce

letters = ascii_lowercase[:11]

def build_multiply_einsten_string(length, ends_with_vector):
    ein_string = ','.join( ''.join(pair) for pair in pairwise(ascii_lowercase[:length+1]))
    return ein_string if not ends_with_vector else ein_string[:-1]

def recursive_matrix_multiply(*args):
    if len(args) == 1:
        return args[0]
    batch = args[:len(letters)-1]
    tail = args[len(letters)-1:]
    return recursive_matrix_multiply( np.einsum(build_multiply_einsten_string(len(batch), batch[-1].ndim == 1 ), *batch), *tail)

# testing ending on a matrix
args = [np.random.rand(5,5) for _ in range(60)]
np.testing.assert_allclose( recursive_matrix_multiply(*args), reduce(np.matmul, args))

# testing ending on a vector
args[-1] = args[-1][:,0]
np.testing.assert_allclose( recursive_matrix_multiply(*args), reduce(np.matmul, args))

请注意,它们适用于两种情况:

  • 纯二维矩阵乘法
  • 终止于向量的点积链
liwlm1x9

liwlm1x92#

我想在Python中高效地执行一系列矩阵向量乘法,numpy.einsum函数似乎是最佳选择。
但最好的选择似乎是numpy.linalg.multi_dot,其中一些文档如下所示:
在单个函数调用中计算两个或多个数组的点积,同时自动选择最快的计算顺序。
multi_dot链接numpy.dot并使用矩阵的最佳括号[1][2]。根据矩阵的形状,这可以大大加快乘法。
如果第一个参数是1-D,则将其视为行向量。如果最后一个参数是1-D,则将其视为列向量。其他参数必须是二维的。
multi_dot视为:

def multi_dot(arrays): return functools.reduce(np.dot, arrays)

这个函数通过简单地将你的矩阵放在一个列表中来解决未知数量的操作的问题。同时,它的性能看起来比einsum更好,一个相对随意的基准测试如下:

In [_]: rng = np.random.default_rng()

In [_]: A = rng.random((10, 20)); B = rng.random((20, 15))

In [_]: C = rng.random((15, 30)); v = rng.random(30)

In [_]: np.allclose(np.linalg.multi_dot([A, B, C, v]),
   ...:             np.einsum('ij,jk,kl,l', A, B, C, v))
Out[_]: True

In [_]: %timeit np.einsum('ij,jk,kl,l', A, B, C, v)
401 µs ± 894 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [_]: %timeit np.linalg.multi_dot([A, B, C, v])
20.3 µs ± 127 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
ncecgwcz

ncecgwcz3#

如注解中所述,您可以使用(未记录?)交错模式,以便于生成下标。就像这样:

np.einsum(*[x for i, a in enumerate(mats) for x in [a, [i, i+1]]], v, [len(mats)]

根据2018年的错误修复,这应该可以达到~50,但是它对我来说不适用于~30。
正如@learning-is-a-mess所建议的,你也可以分块来做:把你的矩阵列表分成N个大小的块,用前一个块的输出作为它的第一个输入,通过einsum运行每个块。
不过,最好的选择可能只是使用opt_einsum。这可以显著提高einsum的速度,并且还支持任意数量的索引:

a = np.array([[1]])
opt_einsum.contract(*[x for i in range(10000) for x in [a, [i, i+1]]])

相关问题