NumPy/SciPy 中的多线程整数矩阵乘法
Multi-threaded integer matrix multiplication in NumPy/SciPy
做类似
的事情
import numpy as np
a = np.random.rand(10**4, 10**4)
b = np.dot(a, a)
使用多核，运行良好。
不过，a 中的元素是 64 位浮点数（或者在 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，但是如果我使用 SciPy BLAS 包装器，即
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 没有的东西'不做。 （此外，b 现在处于 F_CONTIGUOUS 顺序，这是一个较小的问题）。
因此，如果我想进行整数矩阵乘法，我必须执行以下操作之一：
- 使用 NumPy 的速度非常慢
np.dot 很高兴我能保留 8 位整数。
- 使用 SciPy 的
sgemm 并使用 4 倍内存。
- 使用 Numpy 的
np.float16 并且只使用 2 倍内存，需要注意的是 np.dot 在 float16 数组上比在 float32 数组上慢得多，比 int8 慢得多。
- 找到一个用于多线程整数矩阵乘法的优化库（实际上，Mathematica 这样做，但我更喜欢 Python 解决方案），理想情况下支持 1 - 位数组，虽然 8 位数组也很好......（我实际上打算在有限域 Z/2Z 上进行矩阵乘法，我知道我可以用 Sage 做到这一点，这很 Pythonic，但是，再一次，是否有严格的 Python？）
我可以遵循选项 4 吗？有这样的图书馆吗？
免责声明：我实际上是 运行 NumPy + MKL，但我在 vanilly NumPy 上尝试过类似的测试，结果相似。
请注意，虽然这个答案变旧了，但 numpy 可能会获得优化的整数支持。请验证此答案是否仍能在您的设置中更快地运行。
- 选项 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
"" 解释了 为什么 整数这么慢：首先，CPU 有 high-throughput 浮点管道。其次，BLAS没有integer-type.
解决方法： 将矩阵转换为 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
做类似
的事情import numpy as np
a = np.random.rand(10**4, 10**4)
b = np.dot(a, a)
使用多核，运行良好。
不过，a 中的元素是 64 位浮点数（或者在 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，但是如果我使用 SciPy BLAS 包装器，即
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 没有的东西'不做。 （此外，b 现在处于 F_CONTIGUOUS 顺序，这是一个较小的问题）。
因此，如果我想进行整数矩阵乘法，我必须执行以下操作之一：
- 使用 NumPy 的速度非常慢
np.dot很高兴我能保留 8 位整数。 - 使用 SciPy 的
sgemm并使用 4 倍内存。 - 使用 Numpy 的
np.float16并且只使用 2 倍内存，需要注意的是np.dot在 float16 数组上比在 float32 数组上慢得多，比 int8 慢得多。 - 找到一个用于多线程整数矩阵乘法的优化库（实际上，Mathematica 这样做，但我更喜欢 Python 解决方案），理想情况下支持 1 - 位数组，虽然 8 位数组也很好......（我实际上打算在有限域 Z/2Z 上进行矩阵乘法，我知道我可以用 Sage 做到这一点，这很 Pythonic，但是，再一次，是否有严格的 Python？）
我可以遵循选项 4 吗？有这样的图书馆吗？
免责声明：我实际上是 运行 NumPy + MKL，但我在 vanilly NumPy 上尝试过类似的测试，结果相似。
请注意，虽然这个答案变旧了，但 numpy 可能会获得优化的整数支持。请验证此答案是否仍能在您的设置中更快地运行。
- 选项 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
"
解决方法： 将矩阵转换为 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