使用 NumPy 进行快速张量旋转答案

【问题标题】：Fast tensor rotation with NumPy使用 NumPy 进行快速张量旋转
【发布时间】：2011-06-25 03:48:18
【问题描述】：

在应用程序的核心（用 Python 编写并使用 NumPy）我需要旋转一个四阶张量。实际上，我需要多次旋转很多张量，这是我的瓶颈。我的涉及八个嵌套循环的幼稚实现（如下）似乎很慢，但我看不到利用 NumPy 的矩阵运算并希望加快速度的方法。我觉得我应该使用np.tensordot，但我不知道该怎么做。

在数学上，旋转张量 T' 的元素由下式给出： T'_ijkl = Σ g_ia g_jb gkc 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)

有没有办法让它更快？让代码泛化到其他张量等级会很有用，但不太重要。

【问题讨论】：

而且，如果很明显在 numpy 或 scipy 中更快地做到这一点并不容易，我将推出一个 Fortran 扩展模块，看看它是如何执行的。
如果一切都失败了，你可以使用 Cython。据说是plays well with numpy。
虽然我相当确定有一种方法可以减少 numpy 中的嵌套循环（不过我没有立即看到），但正如@delnan 所说，您当前的代码是主要候选者对于 Cython....

标签： python optimization numpy rotation scipy

【解决方案1】：

下面是如何使用单个 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)

诚然，这乍一看有点难以理解，但速度要快很多:)

【讨论】：

+1 表示数学上的出色表现，并再次证明算法优化优于微优化几个数量级。
相当不错！ +1 按照这些思路，未来版本的 numpy 中可能包含一个功能，该功能也符合 OP 的目的......numpy.einsum（还没有）。请参阅关于 numpy-discussion 的讨论：mail-archive.com/numpy-discussion@scipy.org/msg29680.html

【解决方案2】：

要使用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

【讨论】：

btw，cython 版本比第一个代码快 4 倍 stackoverflow.com/questions/4962606/…
发布了close one with four levels of tensordot。您可能会感兴趣的想法:)

【解决方案3】：

出于好奇，我比较了 Cython 实现来自 the question 的幼稚代码与来自 @Philipp's answer 的 numpy 代码。我机器上的 Cython 代码是 four times faster：

#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

【讨论】：

【解决方案4】：

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

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

不过，Philipp 的方法目前快了 3 倍，但也许还有一些改进的余地。速度差异可能主要是由于 tensordot 能够将整个操作展开为可以传递给 BLAS 的单个矩阵乘积，因此避免了与小数组相关的大部分开销——这对于一般的爱因斯坦来说是不可能的求和，因为并非所有可以以这种形式表示的运算都解析为单个矩阵乘积。

【讨论】：

用你的方法来说明my point and approach using four tensordots。 einsum!

【解决方案5】：

我想我会为这些基准测试贡献一个相对较新的数据点，使用 parakeet，这是过去几个月涌现的 numpy 感知 JIT 编译器之一。（我知道的另一个是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

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

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

【讨论】：

我迟到了，但我只是想指出，这些方法随着数据大小的不同而不同。当我将张量从 3x3x3x3 更改为 7x7x7x7 时，Parakeet 和 Numba 在我的机器上都需要大约 10 毫秒，但 Phillip 的第一个解决方案需要大约 80 毫秒。
你试过和the cython implementation of the naive code from the original question比较吗？

【解决方案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)

因此，我们最终得到一个维度张量作为字符串表示法：bcdi。接下来的步骤将涉及将该张量与原始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

【讨论】：

【解决方案7】：

不是一个新的答案，因为以前的所有答案都很好地解决了这个问题。更像是评论，但我将其发布为答案，以便为代码留出一些空间。

虽然所有答案都复制了原始帖子的结果，但我很确定原始帖子中提供的代码是错误的。看公式 T'_ijkl = Σ g_ia g_jb g_kc g_{ld sub> 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 快五倍。

【讨论】：

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