なぜnumpyのeinsumはnumpyの組み込み関数より速いのですか？

質問

の3つの配列から始めましょう。 dtype=np.double . タイミングは、intel CPUで、numpy 1.7.1コンパイルと icc でコンパイルされ、リンクされたintelの mkl . AMDのCPUで、numpy 1.6.1をコンパイルして gcc を使わずに mkl もタイミングを検証するために使用されました。このタイミングはシステムサイズに対してほぼ直線的に変化し、numpy関数で発生する小さなオーバーヘッドによるものではないことに注意してください。 if 文で発生する小さなオーバーヘッドによるものではありません。これらの違いはミリ秒ではなくマイクロ秒で表示されます。

arr_1D=np.arange(500,dtype=np.double)
large_arr_1D=np.arange(100000,dtype=np.double)
arr_2D=np.arange(500**2,dtype=np.double).reshape(500,500)
arr_3D=np.arange(500**3,dtype=np.double).reshape(500,500,500)

まず最初に np.sum 関数を見てみましょう。

np.all(np.sum(arr_3D)==np.einsum('ijk->',arr_3D))
True

%timeit np.sum(arr_3D)
10 loops, best of 3: 142 ms per loop

%timeit np.einsum('ijk->', arr_3D)
10 loops, best of 3: 70.2 ms per loop

パワーズ

np.allclose(arr_3D*arr_3D*arr_3D,np.einsum('ijk,ijk,ijk->ijk',arr_3D,arr_3D,arr_3D))
True

%timeit arr_3D*arr_3D*arr_3D
1 loops, best of 3: 1.32 s per loop

%timeit np.einsum('ijk,ijk,ijk->ijk', arr_3D, arr_3D, arr_3D)
1 loops, best of 3: 694 ms per loop

外装品です。

np.all(np.outer(arr_1D,arr_1D)==np.einsum('i,k->ik',arr_1D,arr_1D))
True

%timeit np.outer(arr_1D, arr_1D)
1000 loops, best of 3: 411 us per loop

%timeit np.einsum('i,k->ik', arr_1D, arr_1D)
1000 loops, best of 3: 245 us per loop

上記はすべて、2 倍の速さで np.einsum . これらの比較は、具体的にはすべて dtype=np.double . 私はこのような操作でスピードアップを期待します。

np.allclose(np.sum(arr_2D*arr_3D),np.einsum('ij,oij->',arr_2D,arr_3D))
True

%timeit np.sum(arr_2D*arr_3D)
1 loops, best of 3: 813 ms per loop

%timeit np.einsum('ij,oij->', arr_2D, arr_3D)
10 loops, best of 3: 85.1 ms per loop

Einsum は、少なくとも 2 倍の速度で np.inner , np.outer , np.kron そして np.sum に関係なく axes を選択することができます。主な例外は np.dot であり、これはBLASライブラリからDGEMMを呼び出しています。では、なぜ np.einsum は同等の他のnumpy関数より速いのでしょうか？

完全性のためにDGEMMの場合。

np.allclose(np.dot(arr_2D,arr_2D),np.einsum('ij,jk',arr_2D,arr_2D))
True

%timeit np.einsum('ij,jk',arr_2D,arr_2D)
10 loops, best of 3: 56.1 ms per loop

%timeit np.dot(arr_2D,arr_2D)
100 loops, best of 3: 5.17 ms per loop

有力な説は@sebergsのコメントからです。 np.einsum を利用することができます。 SSE2 を使用することができますが、numpyのufuncsはnumpy 1.8まで使用できません( 変更履歴 ). 私はこれが正しい答えだと信じています。 ではなく は確認できていません。いくつかの限定的な証拠は、入力配列の d 型を変更し、速度の違いを観察することによって見つけることができ、誰もがタイミングにおける同じ傾向を観察するわけではないという事実があります。

どのように解決するのですか?

ドキュメントによると、すべてのufuncsはSSE2を使用する必要がありますnumpy 1.8がリリースされたので、私はSSE2についてのSebergのコメントが有効であることを再確認したいと思いました。

テストを行うために、新しい python 2.7 インストールを作成し、numpy 1.7 と 1.8 は icc でコンパイルしました。

これは、1.8 アップグレードの前と後の両方のテスト実行です。

import numpy as np
import timeit

arr_1D=np.arange(5000,dtype=np.double)
arr_2D=np.arange(500**2,dtype=np.double).reshape(500,500)
arr_3D=np.arange(500**3,dtype=np.double).reshape(500,500,500)

print 'Summation test:'
print timeit.timeit('np.sum(arr_3D)',
                      'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
                      number=5)/5
print timeit.timeit('np.einsum("ijk->", arr_3D)',
                      'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
                      number=5)/5
print '----------------------\n'


print 'Power test:'
print timeit.timeit('arr_3D*arr_3D*arr_3D',
                      'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
                      number=5)/5
print timeit.timeit('np.einsum("ijk,ijk,ijk->ijk", arr_3D, arr_3D, arr_3D)',
                      'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
                      number=5)/5
print '----------------------\n'


print 'Outer test:'
print timeit.timeit('np.outer(arr_1D, arr_1D)',
                      'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
                      number=5)/5
print timeit.timeit('np.einsum("i,k->ik", arr_1D, arr_1D)',
                      'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
                      number=5)/5
print '----------------------\n'


print 'Einsum test:'
print timeit.timeit('np.sum(arr_2D*arr_3D)',
                      'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
                      number=5)/5
print timeit.timeit('np.einsum("ij,oij->", arr_2D, arr_3D)',
                      'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
                      number=5)/5
print '----------------------\n'

Numpy 1.7.1です。

Summation test:
0.172988510132
0.0934836149216
----------------------

Power test:
1.93524689674
0.839519000053
----------------------

Outer test:
0.130380821228
0.121401786804
----------------------

Einsum test:
0.979052495956
0.126066613197

Numpy 1.8です。

Summation test:
0.116551589966
0.0920487880707
----------------------

Power test:
1.23683619499
0.815982818604
----------------------

Outer test:
0.131808176041
0.127472200394
----------------------

Einsum test:
0.781750011444
0.129271841049

私は、これは、SSE がタイミングの違いに大きな役割を果たしていることをかなり決定的にするものだと思います。 残りの差は、この質問に対する他の回答でカバーされるはずです。