用NumPy快速旋转张量
在应用程序的核心（用Python编写，使用NumPy ）我需要旋转一个四阶张量。 实际上，我需要旋转许多张力，这是我的瓶颈。 我涉及八个嵌套循环的天真实现（下面）似乎很慢，但是我看不到一种方法来利用NumPy的matrix运算，希望能够加快速度。 我有一种感觉，我应该使用np.tensordot ，但我不知道如何。
在math上，旋转张量的元素T'由下式给出：其中，总和在右侧的重复指数之上。 T和Tprime是3 * 3 * 3 * 3的NumPy数组，旋转matrixg是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)
有没有办法让这个更快？ 将代码推广到张量的其他行列将是有用的，但不那么重要。
要使用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的解决scheme快七倍。 如果g张量不经常改变，你也可以cachinggggg张量。 如果你这样做，并打开一些微优化（内联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))
我的家用笔记本电脑上的时间结果（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
下面是如何用一个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)
无可否认，乍看起来有点难以理解，但速度还是比较快的:)
感谢M. Wiebe的努力工作，Numpy的下一个版本（可能会是1.6）将会使这个更容易：
>>> Trot = np.einsum('ai,bj,ck,dl,abcd->ijkl', g, g, g, g, T)
菲利普的做法目前是3倍，但也许还有一些改进空间。 速度的不同可能主要是由于tensordot能够将整个操作展开为一个可传递给BLAS的单个matrix产品，从而避免了与小arrays相关的大部分开销 – 这对于一般的爱因斯坦来说是不可能的总和，因为并非所有可以用这种formsexpression的操作都可以parsing为一个单一的matrix产品。
出于好奇心，我比较了从问题的 Cython实现的朴素代码和@ Philipp的答案中的numpy代码。 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
我想我会为这些基准testing提供一个相对较新的数据点，使用长尾鹦鹉 ，这是过去几个月里出现的具有突出特点的JIT编译器之一。 （另一个我知道的是numba ，但是我没有在这里testing。）
在通过LLVM的一些迷宫般的安装过程之后，您可以装饰许多纯粹的函数以（通常）加速其性能：
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有机会识别和优化热代码path。
有趣的是，斯文和菲利普的build议仍然是数量级更快！
前瞻性方法和解决scheme代码
为了记忆效率和其后的效能，我们可以使用张量matrix乘法。
为了说明所涉及的步骤，让我们使用np.einsum的np.einsum中最简单的解决scheme。 –
np.einsum('ai,bj,ck,dl,abcd->ijkl', g, g, g, g, T)
正如所看到的那样，我们正在通过其四个变体和T之间的张量乘法失去g的第一维。
让我们在步骤中进行张量matrix乘法的和减。 我们从g和T的第一个变体开始：
p1 = np.einsum('abcd, ai->bcdi', T, g)
因此，我们最终得到一个维度的张量作为string符号： bcdi 。 接下来的步骤将涉及总和 – 减less这张张与其他的三个变种，如在原始的einsum使用。 因此，下一个减less将是 –
p2 = np.einsum('bcdi, bj->cdij', p1, g)
如所看到的，我们已经用string符号丢失了前两个维度： a ， b 。 我们继续执行另外两个步骤来摆脱c和d ，并将剩余ijkl作为最终输出，
p3 = np.einsum('cdij, ck->dijk', p2, g) p4 = np.einsum('dijk, dl->ijkl', p3, g)
现在，我们可以使用np.tensordot来减less这些总和，效率会更高。
最终实施
因此，移植到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)))
运行时testing
我们来testing所有在其他post上发布的基于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
正如本文开始时所讨论的，我们正在努力实现内存效率，并因此提升性能。 让我们testing一下，因为我们增加数据集的大小 –
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
