numpy 将np.einsum分解为np.dot操作

xesrikrc  于 2023-10-19  发布在  其他
关注(0)|答案(1)|浏览(100)

我想对四维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的经验有限,所以任何评论都有帮助!
谢谢你,谢谢

1tuwyuhd

1tuwyuhd1#

首先,在numpy中,“Tensor”并不是什么特别的东西。它以相同的方式灵活地处理1d,2d和nd数组。
einsum有一层代码(无论是python还是c),它分析计算,考虑索引和参数的大小,并将其分解为一个或多个计算。在可能的情况下,它将使用与np.dotnp.matmul相同的BLAS库函数。
您可以使用optimize参数和einsum_path函数(请参阅文档)探索其操作。
例如,使用默认的optimize(True),它会将计算“分解”为:

In [166]: print(np.einsum_path('a, b, c, d, abcd', D[i], D[j], D[k], D[l], tensor_S)[1])
  Complete contraction:  a,b,c,d,abcd->
         Naive scaling:  4
     Optimized scaling:  4
      Naive FLOP count:  8.000e+01
  Optimized FLOP count:  6.100e+01
   Theoretical speedup:  1.311
  Largest intermediate:  8.000e+00 elements
--------------------------------------------------------------------------
scaling                  current                                remaining
--------------------------------------------------------------------------
   4                 abcd,a->bcd                              b,c,d,bcd->
   3                   bcd,b->cd                                 c,d,cd->
   2                     cd,c->d                                    d,d->
   1                       d,d->                                       ->

d,d->是一个简单的dot到2 1d数组,产生一个标量。cd,c->d是一个2d和1d,产生一个1d。abcd,a->bcd可以通过将第一个参数重新整形为2d(a,bcd)形状,然后返回到(b,c,d)来完成。np.tensordot进行了这种整形,以将其计算减少到简单的dot
关闭optimizeeinsum使用自己的编译代码一步完成收缩。几年前,我用cythonnditer复制了这种计算,以有效地覆盖整个a,b,c,d空间。

In [167]: print(np.einsum_path('a, b, c, d, abcd', D[i], D[j], D[k], D[l], tensor_S,optimize=False)[1])
  Complete contraction:  a,b,c,d,abcd->
         Naive scaling:  4
     Optimized scaling:  4
      Naive FLOP count:  8.000e+01
  Optimized FLOP count:  8.100e+01
   Theoretical speedup:  0.988
  Largest intermediate:  1.000e+00 elements
--------------------------------------------------------------------------
scaling                  current                                remaining
--------------------------------------------------------------------------
   4              abcd,d,c,b,a->

哪种“优化”选择提供最佳性能可能会有所不同-与索引有关,但也与数组的大小有关。对于您的小示例来说,什么是好的,可能不是大型数组的最佳选择。
经常当SO问到一些复杂的矩阵乘法情况时,我发现用einsum表示最简单。但后来我尝试将其修改为matmul/@(或者在您的情况下是几个)。这可能需要一些转置和/或整形。matmul版本可以更快,但从来没有数量级。
你的einsum可以写成dot的序列,

In [170]: SUM_ijkl
Out[170]: 1325

In [171]: np.dot(D[l], np.dot(D[k], np.dot(D[j], np.dot(D[i], tensor_S.reshape(2,-1)).reshape(2,-1)).reshape(2,-1)).reshape(2,-1))
Out[171]: array([1325])

因为所有的维度都是一样的,所以它不符合我们采取的顺序。如果它们不同,einsum_path可能会选择不同的顺序,最有可能首先压缩最大的,以更快地减少问题空间。
一些比较时间

In [172]: timeit np.dot(D[l], np.dot(D[k], np.dot(D[j], np.dot(D[i], tensor_S.reshape(2,-1)).reshape(2,-1)).reshape(2,-1)).reshape(2,-1))
22.7 µs ± 100 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

In [173]: timeit np.einsum('a, b, c, d, abcd', D[i], D[j], D[k], D[l], tensor_S)
18.3 µs ± 49.7 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

In [174]: timeit np.einsum('a, b, c, d, abcd', D[i], D[j], D[k], D[l], tensor_S, optimize=False)
19.2 µs ± 22.7 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

In [176]: timeit (D[l] @ (D[k] @ (D[j] @ (D[i] @ tensor_S.reshape(2,-1)).reshape(2,-1)).reshape(2,-1)).reshape(2,-1))
22 µs ± 73.4 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

最后一个与np.dot相同,但使用了matmul/@运算符。同样的整形

相关问题