使用NumPy快速进行张量旋转

Question

在一个使用NumPy编写的Python应用程序中，我需要旋转一个4阶张量。实际上，我需要多次旋转许多张量，这是我的瓶颈。我的天真实现（如下）涉及八个嵌套循环，似乎相当慢，但我看不到利用NumPy的矩阵操作并希望加速事情的方法。我有一种感觉我应该使用np.tensordot，但我不知道如何做。

从数学上讲，旋转张量T'的元素由以下公式给出： T'_ijkl = Σ g_ia g_jb g_kc g_ld T_abcd，其中求和是在右手边的重复指数上进行的。 T和Tprime是3*3*3*3 NumPy数组，旋转矩阵g是3*3 NumPy数组。我的缓慢实现（每次调用需要约0.04秒）如下。

#!/usr/bin/env python

import numpy as np

def rotT(T, g):
    Tprime = np.zeros((3,3,3,3))
    for i in range(3):
        for j in range(3):
            for k in range(3):
                for l in range(3):
                    for ii in range(3):
                        for jj in range(3):
                            for kk in range(3):
                                for ll in range(3):
                                    gg = g[ii,i]*g[jj,j]*g[kk,k]*g[ll,l]
                                    Tprime[i,j,k,l] = Tprime[i,j,k,l] + \
                                         gg*T[ii,jj,kk,ll]
    return Tprime

if __name__ == "__main__":

    T = np.array([[[[  4.66533067e+01,  5.84985000e-02, -5.37671310e-01],
                    [  5.84985000e-02,  1.56722231e+01,  2.32831900e-02],
                    [ -5.37671310e-01,  2.32831900e-02,  1.33399259e+01]],
                   [[  4.60051700e-02,  1.54658176e+01,  2.19568200e-02],
                    [  1.54658176e+01, -5.18223500e-02, -1.52814920e-01],
                    [  2.19568200e-02, -1.52814920e-01, -2.43874100e-02]],
                   [[ -5.35577630e-01,  1.95558600e-02,  1.31108757e+01],
                    [  1.95558600e-02, -1.51342210e-01, -6.67615000e-03],
                    [  1.31108757e+01, -6.67615000e-03,  6.90486240e-01]]],
                  [[[  4.60051700e-02,  1.54658176e+01,  2.19568200e-02],
                    [  1.54658176e+01, -5.18223500e-02, -1.52814920e-01],
                    [  2.19568200e-02, -1.52814920e-01, -2.43874100e-02]],
                   [[  1.57414726e+01, -3.86167500e-02, -1.55971950e-01],
                    [ -3.86167500e-02,  4.65601977e+01, -3.57741000e-02],
                    [ -1.55971950e-01, -3.57741000e-02,  1.34215636e+01]],
                   [[  2.58256300e-02, -1.49072770e-01, -7.38843000e-03],
                    [ -1.49072770e-01, -3.63410500e-02,  1.32039847e+01],
                    [ -7.38843000e-03,  1.32039847e+01,  1.38172700e-02]]],
                  [[[ -5.35577630e-01,  1.95558600e-02,  1.31108757e+01],
                    [  1.95558600e-02, -1.51342210e-01, -6.67615000e-03],
                    [  1.31108757e+01, -6.67615000e-03,  6.90486240e-01]],
                   [[  2.58256300e-02, -1.49072770e-01, -7.38843000e-03],
                    [ -1.49072770e-01, -3.63410500e-02,  1.32039847e+01],
                    [ -7.38843000e-03,  1.32039847e+01,  1.38172700e-02]],
                   [[  1.33639532e+01, -1.26331100e-02,  6.84650400e-01],
                    [ -1.26331100e-02,  1.34222177e+01,  1.67851800e-02],
                    [  6.84650400e-01,  1.67851800e-02,  4.89151396e+01]]]])

    g = np.array([[ 0.79389393,  0.54184237,  0.27593346],
                  [-0.59925749,  0.62028664,  0.50609776],
                  [ 0.10306737, -0.56714313,  0.8171449 ]])

    for i in range(100):
        Tprime = rotT(T,g)

有没有办法让它变得更快？使代码普适于其他张量秩可能会很有用，但不是那么重要。

- Andrew Walker

---

如果发现在numpy或scipy中更快地完成这个任务并不容易，我会编写Fortran扩展模块并检查其运行效率。 - Andrew Walker

如果其他方法都失败了，您可以使用Cython。据说它与numpy兼容良好。 - user395760

虽然我相当确定在numpy中有一种方法可以减少嵌套循环（尽管我没有立即看到它），但正如@delnan所说，您当前的代码是Cython的主要候选者... - Joe Kington

Answer 1

要使用tensordot，请计算g张量的外积：

def rotT(T, g):
    gg = np.outer(g, g)
    gggg = np.outer(gg, gg).reshape(4 * g.shape)
    axes = ((0, 2, 4, 6), (0, 1, 2, 3))
    return np.tensordot(gggg, T, axes)

在我的系统上，这比Sven的解决方案快了大约七倍。如果g张量不经常更改，您还可以缓存gggg张量。如果您这样做并打开一些微小的优化（内联tensordot代码，无检查，无通用形状），您仍然可以使其快两倍：

def rotT(T, gggg):
    return np.dot(gggg.transpose((1, 3, 5, 7, 0, 2, 4, 6)).reshape((81, 81)),
                  T.reshape(81, 1)).reshape((3, 3, 3, 3))

我的家用笔记本电脑上使用timeit的结果（500次迭代）：

Your original code: 19.471129179
Sven's code: 0.718412876129
My first code: 0.118047952652
My second code: 0.0690279006958

我的工作电脑上的数字是：

Your original code: 9.77922987938
Sven's code: 0.137110948563
My first code: 0.0569641590118
My second code: 0.0308079719543

- Philipp

---

顺便提一下，Cython版本比第一个代码快4倍。https://dev59.com/Hm445IYBdhLWcg3wQX-G#4973390 - jfs

发布了一个close one with four levels of tensordot。想必对你有兴趣 :) - Divakar

Answer 2

以下是如何用单个Python循环完成此操作的方法：

def rotT(T, g):
    Tprime = T
    for i in range(4):
        slices = [None] * 4
        slices[i] = slice(None)
        slices *= 2
        Tprime = g[slices].T * Tprime
    return Tprime.sum(-1).sum(-1).sum(-1).sum(-1)

不可否认，一开始可能有点难以理解，但速度会更快 :)

- Sven Marnach

---

数学的精妙和算法优化的胜利再次证明，相较微观优化而言，算法优化的效果更为显著。+1 - user395760

相当不错！+1 顺着这个思路，未来版本的numpy可能会包括一个函数，也适合OP的目的...numpy.einsum（目前还没有）。请参阅numpy-discussion上的此讨论：http://www.mail-archive.com/[e-mail]/msg29680.html - Joe Kington

Answer 3

多亏了M. Wiebe的辛勤工作，下一个版本的Numpy（可能是1.6）将使这更加容易：

>>> Trot = np.einsum('ai,bj,ck,dl,abcd->ijkl', g, g, g, g, T)

目前，Philipp的方法快了3倍，但也许还有改进的空间。速度差异可能主要是由于tensordot能够将整个操作展开为单个矩阵乘积并传递给BLAS，从而避免了与小数组相关的大量开销 --- 对于一般的Einstein求和来说，这是不可能的，因为并非所有可以用这种形式表达的操作都会解析为单个矩阵乘积。

- pv.

---

使用你的方法来说明我的观点和方法，使用四个tensordots。用einsum简洁明了！ - Divakar

Answer 4

出于好奇，我比较了Cython实现的一个天真代码和问题中的numpy代码以及@Philipp's answer。在我的机器上，Cython代码的速度是原始代码的四倍：

#cython: boundscheck=False, wraparound=False
import numpy as np
cimport numpy as np

def rotT(np.ndarray[np.float64_t, ndim=4] T,
         np.ndarray[np.float64_t, ndim=2] g):
    cdef np.ndarray[np.float64_t, ndim=4] Tprime
    cdef Py_ssize_t i, j, k, l, ii, jj, kk, ll
    cdef np.float64_t gg

    Tprime = np.zeros((3,3,3,3), dtype=T.dtype)
    for i in range(3):
        for j in range(3):
            for k in range(3):
                for l in range(3):
                    for ii in range(3):
                        for jj in range(3):
                            for kk in range(3):
                                for ll in range(3):
                                    gg = g[ii,i]*g[jj,j]*g[kk,k]*g[ll,l]
                                    Tprime[i,j,k,l] = Tprime[i,j,k,l] + \
                                         gg*T[ii,jj,kk,ll]
    return Tprime

- jfs

Answer 5

我想为这些基准测试做出一个相对较新的数据点，使用了最近几个月中出现的numpy感知JIT编译器之一——parakeet。我知道的另一个是numba，但我没有在这里测试它。您可以通过安装LLVM并对许多纯numpy函数进行修饰来（通常）加速它们的性能，即使安装过程有些复杂。

import numpy as np
import parakeet

@parakeet.jit
def rotT(T, g):
    # ...

我只尝试将JIT应用于原始问题中的Andrew代码，但它表现得相当不错（速度提高了10倍以上），而无需编写任何新代码：

andrew      10 loops, best of 3: 206 msec per loop
andrew_jit  10 loops, best of 3: 13.3 msec per loop
sven        100 loops, best of 3: 2.39 msec per loop
philipp     1000 loops, best of 3: 0.879 msec per loop

对于这些时间（在我的笔记本电脑上），我运行了每个函数十次，以便JIT有机会识别和优化热代码路径。

有趣的是，Sven和Philipp的建议仍然快几个数量级！

- lmjohns3

---

我加入这段对话有些晚，但我想指出这些方法对数据大小的扩展不同。当我将张量从3x3x3x3改变为7x7x7x7时，在我的机器上Parakeet和Numba都需要约10毫秒，但Phillip的第一个解决方案需要约80毫秒。 - Alex Rubinsteyn

你尝试过将其与原始问题中的Cython实现的朴素代码进行比较吗？ - jfs

Answer 6

前瞻性方法和解决方案代码

为了提高内存效率和性能效率，我们可以分步使用张量矩阵乘法。

为了说明所涉及的步骤，让我们使用最简单的解决方案之一，即np.einsum by @pv。 -

np.einsum('ai,bj,ck,dl,abcd->ijkl', g, g, g, g, T)

如上所述，通过将g的四个变量与T进行张量乘法，我们正在失去第一维。

让我们逐步完成张量矩阵乘法中的求和约简。让我们从g和T的第一个变量开始：

p1 = np.einsum('abcd, ai->bcdi', T, g)

因此，我们最终得到一个维度为字符串符号的张量。接下来的步骤将涉及将此张量与原始einsum实现中使用的其余三个g变量进行求和约简。因此，下一步的约简将是 -

p2 = np.einsum('bcdi, bj->cdij', p1, g)

正如所见，我们使用字符串标记a、b失去了前两个维度。我们继续进行两个步骤，以摆脱c和d，最终输出为ijkl，如下所示 -

p3 = np.einsum('cdij, ck->dijk', p2, g)

p4 = np.einsum('dijk, dl->ijkl', p3, g)

现在，我们可以使用np.tensordot来进行这些求和操作，这将更加高效。

最终实现

因此，转移到np.tensordot后，我们的最终实现如下 -

p1 = np.tensordot(T,g,axes=((0),(0)))
p2 = np.tensordot(p1,g,axes=((0),(0)))
p3 = np.tensordot(p2,g,axes=((0),(0)))
out = np.tensordot(p3,g,axes=((0),(0)))

运行时测试

让我们测试其他帖子中发布的基于NumPy的方法来解决性能问题。

作为函数的方法 -

def rotT_Philipp(T, g):  # @Philipp's soln
    gg = np.outer(g, g)
    gggg = np.outer(gg, gg).reshape(4 * g.shape)
    axes = ((0, 2, 4, 6), (0, 1, 2, 3))
    return np.tensordot(gggg, T, axes)

def rotT_Sven(T, g):    # @Sven Marnach's soln
    Tprime = T
    for i in range(4):
        slices = [None] * 4
        slices[i] = slice(None)
        slices *= 2
        Tprime = g[slices].T * Tprime
    return Tprime.sum(-1).sum(-1).sum(-1).sum(-1)    

def rotT_pv(T, g):     # @pv.'s soln
    return np.einsum('ai,bj,ck,dl,abcd->ijkl', g, g, g, g, T)

def rotT_Divakar(T,g): # Posted in this post
    p1 = np.tensordot(T,g,axes=((0),(0)))
    p2 = np.tensordot(p1,g,axes=((0),(0)))
    p3 = np.tensordot(p2,g,axes=((0),(0)))
    p4 = np.tensordot(p3,g,axes=((0),(0)))
    return p4

使用原始数据集大小的时间 -

In [304]: # Setup inputs 
     ...: T = np.random.rand(3,3,3,3)
     ...: g = np.random.rand(3,3)
     ...: 

In [305]: %timeit rotT(T, g)
     ...: %timeit rotT_pv(T, g)
     ...: %timeit rotT_Sven(T, g)
     ...: %timeit rotT_Philipp(T, g)
     ...: %timeit rotT_Divakar(T, g)
     ...: 
100 loops, best of 3: 6.51 ms per loop
1000 loops, best of 3: 247 µs per loop
10000 loops, best of 3: 137 µs per loop
10000 loops, best of 3: 41.6 µs per loop
10000 loops, best of 3: 28.3 µs per loop

In [306]: 6510.0/28.3 # Speedup with the proposed soln over original code
Out[306]: 230.03533568904592

如本文开头所讨论，我们试图通过提高内存效率来提升性能。随着数据集的增大，让我们测试一下这一点 -

In [307]: # Setup inputs 
     ...: T = np.random.rand(5,5,5,5)
     ...: g = np.random.rand(5,5)
     ...: 

In [308]: %timeit rotT(T, g)
     ...: %timeit rotT_pv(T, g)
     ...: %timeit rotT_Sven(T, g)
     ...: %timeit rotT_Philipp(T, g)
     ...: %timeit rotT_Divakar(T, g)
     ...: 
100 loops, best of 3: 6.54 ms per loop
100 loops, best of 3: 7.17 ms per loop
100 loops, best of 3: 2.7 ms per loop
1000 loops, best of 3: 1.47 ms per loop
10000 loops, best of 3: 39.9 µs per loop

- Divakar

Answer 7

并不是一个新的答案，因为之前所有的回答都很好地回答了这个问题。更像是一条评论，但我将它发布为一个答案来为代码留些空间。

虽然所有的回答都复现了原始帖子的结果，但我非常确定原始帖子中提供的代码是错误的。看着公式T'_ijkl = Σ g_ia g_jb g_kc g_ld T_abcd，我相信它是正确的，计算每个T'条目时变化的g指数是a、b、c和d。然而，在原始提供的代码中，用于访问计算gg值的g值的索引与公式相比被交换了。因此，我相信以下代码实际上提供了正确的公式实现：

def rotT(T, g):
    Tprime = np.zeros((3, 3, 3, 3))
    for i in range(3):
        for j in range(3):
            for k in range(3):
                for l in range(3):
                    for a in range(3):
                        for b in range(3):
                            for c in range(3):
                                for d in range(3):
                                    Tprime[i, j, k, l] += \
                                        g[i, a] * g[j, b] * \
                                        g[k, c] * g[l, d] * T[a, b, c, d]

等效但更快的einsum和tensordot调用更新为：

Tprime = np.tensordot(g, np.tensordot(g, np.tensordot(
    g, np.tensordot(g, T, (1, 3)), (1, 3)), (1, 3)), (1, 3))
Tprime = np.einsum('ia, jb, kc, ld, abcd->ijkl', g, g, g, g, T)

此外，在我的计算机上，对于朴素循环函数，使用 Numba 的 @jit(nopython=True) 比使用 numpy.tensordot 快五倍。

- Patol75

---

很高兴看到Numba终于在这个帖子中出现了。我认为没有人提到Python for循环很慢的基本事实，但是如果使用Numba编译它们，所有这些循环都会转换为运行非常快的LVLL代码。 Numba允许您利用并行操作，甚至是GPU，如果您真的需要速度。 - Diomedea