我想对四维Tensor和相应的系数求和。可以使用4个嵌套的for
循环来获得正确的答案。
但是我们都知道for
循环在Python中是两个慢的。这里我使用np.einsum来得到一个结果,比如:
import numpy as np
D = np.arange(4).reshape(2, 2)
i, j, k, l = 1, 0, 1, 1 # index < 2
tensor_S = np.arange(16).reshape(2, 2, 2, 2)
SUM_ijkl = np.einsum('a, b, c, d, abcd', D[i], D[j], D[k], D[l], tensor_S)
这看起来很简洁,但我了解到np.einsum
是用纯Python实现的,请参阅here,所以当遇到大型Tensor时,这可能是一个效率问题。我想知道如何将这种求和分解为点积或矩阵乘法来保证计算效率,或者通过其他替代的快速函数得到结果。
我处理Tensor的经验有限,所以任何评论都有帮助!
谢谢你,谢谢
1条答案
按热度按时间1tuwyuhd1#
首先,在
numpy
中,“Tensor”并不是什么特别的东西。它以相同的方式灵活地处理1d,2d和nd数组。einsum
有一层代码(无论是python还是c),它分析计算,考虑索引和参数的大小,并将其分解为一个或多个计算。在可能的情况下,它将使用与np.dot
和np.matmul
相同的BLAS库函数。您可以使用
optimize
参数和einsum_path
函数(请参阅文档)探索其操作。例如,使用默认的optimize(
True
),它会将计算“分解”为:d,d->
是一个简单的dot
到2 1d数组,产生一个标量。cd,c->d
是一个2d和1d,产生一个1d。abcd,a->bcd
可以通过将第一个参数重新整形为2d(a,bcd)形状,然后返回到(b,c,d)来完成。np.tensordot
进行了这种整形,以将其计算减少到简单的dot
。关闭
optimize
,einsum
使用自己的编译代码一步完成收缩。几年前,我用cython
和nditer
复制了这种计算,以有效地覆盖整个a,b,c,d
空间。哪种“优化”选择提供最佳性能可能会有所不同-与索引有关,但也与数组的大小有关。对于您的小示例来说,什么是好的,可能不是大型数组的最佳选择。
经常当SO问到一些复杂的矩阵乘法情况时,我发现用
einsum
表示最简单。但后来我尝试将其修改为matmul/@
(或者在您的情况下是几个)。这可能需要一些转置和/或整形。matmul
版本可以更快,但从来没有数量级。你的
einsum
可以写成dot
的序列,因为所有的维度都是一样的,所以它不符合我们采取的顺序。如果它们不同,
einsum_path
可能会选择不同的顺序,最有可能首先压缩最大的,以更快地减少问题空间。一些比较时间
最后一个与
np.dot
相同,但使用了matmul/@
运算符。同样的整形