并行化正在执行位操作的代码

html5 • • 问答

我有这段代码通过将这个 AU 矩阵的每字节 8 个元素打包到 A 中来减少内存消耗，从而使 100k*200k 矩阵占用更少的空间。如您所料，此代码需要永远运行，我还计划将行数增加到 200k。我在一个非常强大的实例（CPU 和 GPU）上运行代码并且可以扩展它，所以任何人都可以帮助并行化此代码，使其更快。

import numpy as np
colm = int(2000000/8)
rows = 1000000
cols = int(colm*8)
AU = np.random.randint(2,size=(rows, cols),dtype=np.int8)
start_time = time.time()

A = np.empty((rows,colm), dtype=np.uint8)
for i in range(A.shape[0]):
    for j in range(A.shape[1]):
        A[i,j] = 0
        for k in range(8):
            if AU[i,(j*8)+k] == 1:
                A[i,j] = A[i,j] | (1<<(7-k))

回答

警告：您尝试分配大量内存：您可能没有大约2 TB 的内存。

假设你有足够的内存或者你可以减少数据集的大小，你可以使用Numba JIT编写一个更快的实现。此外，您可以并行化代码并用无分支实现替换慢速条件，以显着加快计算速度，因为AU它填充了二进制值。最后，您可以展开正在处理的内部循环k以使代码更快。这是由此产生的实现：

import numpy as np
import numba as nb
colm = int(2000000/8)
rows = 1000000
cols = int(colm*8)
AU = np.random.randint(2,size=(rows, cols),dtype=np.int8)
A = np.empty((rows,colm), dtype=np.uint8)

@nb.njit('void(uint8[:,:],int8[:,:])', parallel=True)
def compute(A, AU):
    for i in nb.prange(A.shape[0]):
        for j in range(A.shape[1]):
            offset = j * 8
            res = AU[i,offset] << 7
            res |= AU[i,offset+1] << 6
            res |= AU[i,offset+2] << 5
            res |= AU[i,offset+3] << 4
            res |= AU[i,offset+4] << 3
            res |= AU[i,offset+5] << 2
            res |= AU[i,offset+6] << 1
            res |= AU[i,offset+7]
            A[i,j] = res

compute(A, AU)

在我的机器上，此代码比在较小数据集（带有和）上的原始实现快37851倍。最初的实现需要 6 分钟 3 秒，而优化的需要 9.6 毫秒。colm=int(20000/8)rows=10000

这段代码在我的机器上受内存限制。对于当前输入，此代码接近最优，因为它大部分时间都在读取AU输入矩阵。一个好的额外优化是将AU矩阵“压缩”为更小的矩阵（如果可能）。

以上是并行化正在执行位操作的代码的全部内容。
THE END
分享
二维码
< <上一篇
下一篇>>