python style 使用NumPy進行快速張量旋轉




seaborn vs matplotlib (5)

應用程序的核心(用Python編寫並使用NumPy )我需要旋轉四階張量。 實際上,我需要多次旋轉很多張量,這是我的瓶頸。 我的天真實現(下面)涉及八個嵌套循環似乎相當慢,但我看不到一種方法來利用NumPy的矩陣運算,並希望加快速度。 我有一種感覺,我應該使用np.tensordot ,但我不知道如何。

在數學上,旋轉張量的元素T'由下式給出: T'ijjl = Σgia 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)

有沒有辦法讓這個更快? 將代碼推廣到張量的其他等級將是有用的,但不太重要。


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

感謝M. Wiebe的辛勤工作,Numpy的下一個版本(可能是1.6)將使這更容易:

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

不過,菲利普的方法目前要快3倍,但也許還有一些改進空間。 速度差異可能主要是由於tensordot能夠將整個操作展開為可以傳遞給BLAS的單個矩陣產品,因此避免了與小陣列相關的大量開銷 - 這對於一般的愛因斯坦來說是不可能的求和,因為並非所有可以用這種形式表達的操作都能解析為單個矩陣產品。


前瞻性方法和解決方案代碼

為了提高內存效率和性能效率,我們可以逐步使用張量矩陣乘法。

為了說明所涉及的步驟,讓我們使用np.einsum最簡單的np.einsum解決方案。 -

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

如圖所示,我們正在失去g的第一個維度,其四個變量和T之間存在張量乘法。

讓我們逐步減少張量矩陣乘法的和。 讓我們從gT的第一個變體開始:

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

因此,我們最終得到一個尺寸的張量作為字符串表示法: bcdi 。 接下來的步驟將涉及總和 - 將該張量與原始的einsum實施中使用的三個g變體的其餘部分相比減少。 因此,下一次減少將是 -

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

如圖所示,我們使用字符串符號丟失了前兩個維度: ab 。 我們繼續兩個步驟來擺脫cd並留下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

我認為我會使用parakeet為這些基準提供一個相對較新的數據點, parakeet是過去幾個月湧現出來的numpy -aware JIT編譯器之一。 (另一個我知道的是numba ,但我沒有在這裡測試它。)

在通過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有機會識別和優化熱代碼路徑。

有趣的是,斯文和菲利普的建議仍然快了幾個數量級!


以下是使用單個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)

不可否認,乍一看這有點難以理解,但它的速度要快得多:)





scipy