web-dev-qa-db-ja.com

NumPy / SciPyでのマルチスレッド整数行列乗算

次のようなことをする

import numpy as np
a = np.random.Rand(10**4, 10**4)
b = np.dot(a, a)

複数のコアを使用し、うまく動作します。

ただし、aの要素は64ビットのfloat(または、32ビットプラットフォームでは32ビット?)であり、8ビットの整数配列を乗算したいと思います。ただし、次のことを試してください。

a = np.random.randint(2, size=(n, n)).astype(np.int8)

その結果、ドット積は複数のコアを使用しないため、PCでの実行速度が約1000倍遅くなります。

array: np.random.randint(2, size=shape).astype(dtype)

dtype    shape          %time (average)

float32 (2000, 2000)    62.5 ms
float32 (3000, 3000)    219 ms
float32 (4000, 4000)    328 ms
float32 (10000, 10000)  4.09 s

int8    (2000, 2000)    13 seconds
int8    (3000, 3000)    3min 26s
int8    (4000, 4000)    12min 20s
int8    (10000, 10000)  It didn't finish in 6 hours

float16 (2000, 2000)    2min 25s
float16 (3000, 3000)    Not tested
float16 (4000, 4000)    Not tested
float16 (10000, 10000)  Not tested

NumPyが整数をサポートしないBLASを使用していることは理解していますが、SciPyBLASラッパーを使用している場合は次のようになります。

import scipy.linalg.blas as blas
a = np.random.randint(2, size=(n, n)).astype(np.int8)
b = blas.sgemm(alpha=1.0, a=a, b=a)

計算マルチスレッドです。現在、blas.sgemmはfloat32の場合はnp.dotとまったく同じタイミングで実行されますが、非浮動小数点の場合はすべてをfloat32に変換し、浮動小数点数を出力します。これはnp.dotではありません。します。 (さらに、bF_CONTIGUOUSの順序になりましたが、これはそれほど問題ではありません)。

したがって、整数行列の乗算を実行する場合は、次のいずれかを実行する必要があります。

  1. NumPyの痛々しいほど遅いnp.dotを使用して、8ビット整数を維持できてうれしいです。
  2. SciPyのsgemmを使用し、4倍のメモリを使い果たします。
  3. Numpyのnp.float16を使用し、2xメモリのみを使用します。ただし、np.dotはfloat32配列よりもfloat16配列の方がはるかに遅く、int8よりも遅くなります。
  4. マルチスレッド整数行列乗算用に最適化されたライブラリを見つけます(実際には、Mathematicaはこれを行いますが、理想的には1ビット配列をサポートするPythonソリューション)をお勧めします、8ビット配列も問題ありませんが...(私は実際には有限フィールドZ/2Zで行列の乗算を行うことを目指しており、これはSageで実行できることを知っています。かなりPythonicですが、繰り返しになりますが、厳密にPythonの何かがありますか?)

オプション4に従うことはできますか?そのようなライブラリは存在しますか?

免責事項:私は実際にNumPy + MKLを実行していますが、バニリーNumPyで同様のテストを試しましたが、同様の結果が得られました。

24
  • オプション5-カスタムソリューションをロールします:行列積をいくつかのサブ積に分割し、これらを並行して実行します。これは、標準のPythonモジュールで比較的簡単に実装できます。サブプロダクトは、グローバルインタープリターロックを解放するnumpy.dotで計算されます。したがって、 threads は比較的軽量で、メモリ効率のためにメインスレッドから配列にアクセスできます。

実装:

import numpy as np
from numpy.testing import assert_array_equal
import threading
from time import time


def blockshaped(arr, nrows, ncols):
    """
    Return an array of shape (nrows, ncols, n, m) where
    n * nrows, m * ncols = arr.shape.
    This should be a view of the original array.
    """
    h, w = arr.shape
    n, m = h // nrows, w // ncols
    return arr.reshape(nrows, n, ncols, m).swapaxes(1, 2)


def do_dot(a, b, out):
    #np.dot(a, b, out)  # does not work. maybe because out is not C-contiguous?
    out[:] = np.dot(a, b)  # less efficient because the output is stored in a temporary array?


def pardot(a, b, nblocks, mblocks, dot_func=do_dot):
    """
    Return the matrix product a * b.
    The product is split into nblocks * mblocks partitions that are performed
    in parallel threads.
    """
    n_jobs = nblocks * mblocks
    print('running {} jobs in parallel'.format(n_jobs))

    out = np.empty((a.shape[0], b.shape[1]), dtype=a.dtype)

    out_blocks = blockshaped(out, nblocks, mblocks)
    a_blocks = blockshaped(a, nblocks, 1)
    b_blocks = blockshaped(b, 1, mblocks)

    threads = []
    for i in range(nblocks):
        for j in range(mblocks):
            th = threading.Thread(target=dot_func, 
                                  args=(a_blocks[i, 0, :, :], 
                                        b_blocks[0, j, :, :], 
                                        out_blocks[i, j, :, :]))
            th.start()
            threads.append(th)

    for th in threads:
        th.join()

    return out


if __name__ == '__main__':
    a = np.ones((4, 3), dtype=int)
    b = np.arange(18, dtype=int).reshape(3, 6)
    assert_array_equal(pardot(a, b, 2, 2), np.dot(a, b))

    a = np.random.randn(1500, 1500).astype(int)

    start = time()
    pardot(a, a, 2, 4)
    time_par = time() - start
    print('pardot: {:.2f} seconds taken'.format(time_par))

    start = time()
    np.dot(a, a)
    time_dot = time() - start
    print('np.dot: {:.2f} seconds taken'.format(time_dot))

この実装により、マシンのコアの物理数である約x4のスピードアップが得られます。

running 8 jobs in parallel
pardot: 5.45 seconds taken
np.dot: 22.30 seconds taken
5
kazemakase

" int by intと比較してfloatby float行列の乗算を実行する方が速いのはなぜですか? "説明なぜ整数が非常に遅いのか:まず、CPUは高スループットの浮動小数点パイプライン。第二に、BLASには整数型がありません。

回避策:行列をfloat32値に変換すると、大幅に高速化されます。 2015 MacBook Proの90倍のスピードアップはどうですか? (float64の使用は半分です。)

import numpy as np
import time

def timeit(callable):
    start = time.time()
    callable()
    end = time.time()
    return end - start

a = np.random.random_integers(0, 9, size=(1000, 1000)).astype(np.int8)

timeit(lambda: a.dot(a))  # ≈0.9 sec
timeit(lambda: a.astype(np.float32).dot(a.astype(np.float32)).astype(np.int8) )  # ≈0.01 sec
1
Jerry101