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

Question

dtype=np.doubleの3つの配列から始めましょう。タイミングは、iccでコンパイルされ、インテルのmklにリンクされたnumpy 1.7.1を使用して、インテルCPUで実行されます。 gccなしでmklでコンパイルされたnumpy 1.6.1を搭載したAMD CPUも、タイミングの検証に使用されました。タイミングはシステムサイズにほぼ比例してスケールし、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

上記のすべては、np.einsumで2倍高速です。これらはすべて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は、axesの選択に関係なく、np.inner、np.outer、np.kron、およびnp.sumの少なくとも2倍の速度であるようです。 BLASライブラリからDGEMMを呼び出すときの主な例外はnp.dotです。では、なぜ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

---

主要な理論は、np.einsumが SSE2 を使用できるという@sebergsコメントからですが、numpyのufuncsはnumpy 1.8までです（ change log を参照）。これは正しい答えだと思いますが、not確認できませんでした。入力配列のdtypeを変更し、速度の違いと、すべての人がタイミングの同じ傾向を観察するわけではないという事実を観察することにより、いくつかの限定的な証拠を見つけることができます。

Answer 1

Numpy 1.8がリリースされたので、ドキュメントによると、すべてのufuncsがSSE2を使用するはずであるため、SSE2に関するSebergのコメントが有効であることを再確認したかったのです。

テストを実行するために、新しいpython 2.7インストールが作成されました。numpy1.7および1.8は、Ubuntuを実行しているAMD opteronコアの標準オプションを使用して、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はタイミングの違いに大きな役割を果たしているという結論に達していると思います。この質問に対する他の回答。

Answer 2

まず、numpyリストでこれについての過去の議論がたくさんありました。たとえば、次を参照してください： http://numpy-discussion.10968.n7.nabble.com/poor-performance-of-sum-with-sub-machine-Word-integer-types-td41.html = http://numpy-discussion.10968.n7.nabble.com/odd-performance-of-sum-td3332.html

いくつかは、einsumが新しく、キャッシュのアライメントやその他のメモリアクセスの問題について改善しようとしているという事実に要約されています。一方、古いnumpy関数の多くは、非常に最適化された1。ただ、推測しているところです。

---

ただし、あなたがしていることのいくつかは、かなり「リンゴ対リンゴ」の比較ではありません。

@Jamieがすでに言ったことに加えて、sumは配列により適切なアキュムレータを使用します

たとえば、sumは、入力のタイプをチェックし、適切なアキュムレーターを使用することにより注意を払っています。たとえば、次のことを考慮してください。

In [1]: x = 255 * np.ones(100, dtype=np.uint8)

In [2]: x
Out[2]:
array([255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
       255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
       255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
       255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
       255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
       255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
       255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
       255, 255, 255, 255, 255, 255, 255, 255, 255], dtype=uint8)

sumが正しいことに注意してください。

In [3]: x.sum()
Out[3]: 25500

einsumは間違った結果を返しますが：

In [4]: np.einsum('i->', x)
Out[4]: 156

しかし、あまり制限のないdtypeを使用しても、期待どおりの結果が得られます。

In [5]: y = 255 * np.ones(100)

In [6]: np.einsum('i->', y)
Out[6]: 25500.0

Answer 3

私はこれらのタイミングが何が起こっているのかを説明すると思う：

a = np.arange(1000, dtype=np.double)
%timeit np.einsum('i->', a)
100000 loops, best of 3: 3.32 us per loop
%timeit np.sum(a)
100000 loops, best of 3: 6.84 us per loop

a = np.arange(10000, dtype=np.double)
%timeit np.einsum('i->', a)
100000 loops, best of 3: 12.6 us per loop
%timeit np.sum(a)
100000 loops, best of 3: 16.5 us per loop

a = np.arange(100000, dtype=np.double)
%timeit np.einsum('i->', a)
10000 loops, best of 3: 103 us per loop
%timeit np.sum(a)
10000 loops, best of 3: 109 us per loop

したがって、np.sumをnp.einsum経由で呼び出すと、基本的に3usのオーバーヘッドがほぼ一定になるため、基本的には高速に実行されますが、処理に少し時間がかかります。なぜそうなるのでしょうか？私のお金は次のとおりです。

a = np.arange(1000, dtype=object)
%timeit np.einsum('i->', a)
Traceback (most recent call last):
...
TypeError: invalid data type for einsum
%timeit np.sum(a)
10000 loops, best of 3: 20.3 us per loop

正確に何が起こっているのかわかりませんが、np.einsumは乗算と加算を行うために型固有の関数を抽出するためにいくつかのチェックをスキップしており、*と+標準Cタイプのみ。

---

多次元の場合は違いません：

n = 10; a = np.arange(n**3, dtype=np.double).reshape(n, n, n)
%timeit np.einsum('ijk->', a)
100000 loops, best of 3: 3.79 us per loop
%timeit np.sum(a)
100000 loops, best of 3: 7.33 us per loop

n = 100; a = np.arange(n**3, dtype=np.double).reshape(n, n, n)
%timeit np.einsum('ijk->', a)
1000 loops, best of 3: 1.2 ms per loop
%timeit np.sum(a)
1000 loops, best of 3: 1.23 ms per loop

そのため、ほとんど一定のオーバーヘッドが発生しますが、いったんそれが達成されると高速に実行されることはありません。

Answer 4

Numpy 1.16.4の更新：Numpyのネイティブ関数は、ほとんどすべての場合でeinsumsよりも高速です。 einsumの外部バリアントとsum23のみが、非einsumバージョンよりも高速にテストされます。

Numpyのネイティブ関数を使用できる場合は、それを行います。

（私のプロジェクトである perfplot で作成された画像。）

---

プロットを再現するコード：

import numpy
import perfplot


def setup1(n):
    return numpy.arange(n, dtype=numpy.double)


def setup2(n):
    return numpy.arange(n ** 2, dtype=numpy.double).reshape(n, n)


def setup3(n):
    return numpy.arange(n ** 3, dtype=numpy.double).reshape(n, n, n)

def setup23(n):
    return (
        numpy.arange(n ** 2, dtype=numpy.double).reshape(n, n),
        numpy.arange(n ** 3, dtype=numpy.double).reshape(n, n, n)
    )


def numpy_sum(a):
    return numpy.sum(a)


def einsum_sum(a):
    return numpy.einsum("ijk->", a)


perfplot.save(
    "sum.png",
    setup=setup3,
    kernels=[numpy_sum, einsum_sum],
    n_range=[2 ** k for k in range(10)],
    logx=True,
    logy=True,
    title="sum",
)


def numpy_power(a):
    return a * a * a


def einsum_power(a):
    return numpy.einsum("ijk,ijk,ijk->ijk", a, a, a)


perfplot.save(
    "power.png",
    setup=setup3,
    kernels=[numpy_power, einsum_power],
    n_range=[2 ** k for k in range(9)],
    logx=True,
    logy=True,
)


def numpy_outer(a):
    return numpy.outer(a, a)


def einsum_outer(a):
    return numpy.einsum("i,k->ik", a, a)


perfplot.save(
    "outer.png",
    setup=setup1,
    kernels=[numpy_outer, einsum_outer],
    n_range=[2 ** k for k in range(13)],
    logx=True,
    logy=True,
    title="outer",
)



def dgemm_numpy(a):
    return numpy.dot(a, a)


def dgemm_einsum(a):
    return numpy.einsum("ij,jk", a, a)


def dgemm_einsum_optimize(a):
    return numpy.einsum("ij,jk", a, a, optimize=True)


perfplot.save(
    "dgemm.png",
    setup=setup2,
    kernels=[dgemm_numpy, dgemm_einsum],
    n_range=[2 ** k for k in range(13)],
    logx=True,
    logy=True,
    title="dgemm",
)



def dot_numpy(a):
    return numpy.dot(a, a)


def dot_einsum(a):
    return numpy.einsum("i,i->", a, a)


perfplot.save(
    "dot.png",
    setup=setup1,
    kernels=[dot_numpy, dot_einsum],
    n_range=[2 ** k for k in range(20)],
    logx=True,
    logy=True,
    title="dot",
)

def sum23_numpy(data):
    a, b = data
    return numpy.sum(a * b)


def sum23_einsum(data):
    a, b = data
    return numpy.einsum('ij,oij->', a, b)


perfplot.save(
    "sum23.png",
    setup=setup23,
    kernels=[sum23_numpy, sum23_einsum],
    n_range=[2 ** k for k in range(10)],
    logx=True,
    logy=True,
    title="sum23",
)