对两个大矩阵执行矩阵运算

Question

我有两个大矩阵 (40000*4096)，我想将第一个矩阵的每一行与第二个矩阵的所有行进行比较和匹配，因此，输出的大小为 (40000* 40000）。但是，由于我需要这样做数千次，因此每次迭代需要花费 26k 秒的时间，所以 5000 次...... 如果你能给我一些聪明的建议，我会很高兴。谢谢你。 附言到目前为止，我只做了一次迭代（5000 次中的 1 次）

def matcher(Antigens, Antibodies,ind):
    temp = np.zeros((Antibodies.shape[0],Antibodies.shape[1]))
    output = np.zeros((Antibodies.shape[0],1))
    for i in range(len(Antibodies)):
        temp[i] = np.int32(np.equal(Antigens[ind],Antibodies[i]))
        output[i] = np.sum(temp[i])
    return output
output = [matcher(gens,Antibodies) for gens in Antigens]

Answer 1

好的，我想我明白你的目标是什么了：

计算行匹配数（抗原与抗体矩阵）。结果向量的每一行 (40,000 x 1) 代表 1 个抗原行和所有抗体行之间完全匹配的计数（因此值从 0 到 40_000）。

我做了一些假数据：

import numpy as np
import numba as nb

num_mat = 5       # number of matrices
num_row = 10_000  # number of rows per matrix
num_elm = 4_096   # number of elements per row
dim = (num_mat,num_row,num_elm)

Antigens = np.random.randint(0,256,dim,dtype=np.uint8)
Antibodies = np.random.randint(0,256,dim,dtype=np.uint8)

这里有一点很重要，我将矩阵简化为可以表示数据的最小数据类型，以减少它们的内存占用。我不确定您的数据是什么样的，但希望您也可以这样做。

此外，以下代码假设您的尺寸看起来像假数据：

（矩阵、行、元素的数量）

@nb.njit
def match_arr(arr1, arr2):
    for i in range(arr1.shape[0]): #4096 vs 4096
        if arr1[i] != arr2[i]:
            return False
    return True

@nb.njit
def match_mat_sum(ag, ab):
    out = np.zeros((ag.shape[0])) # 40000
    for i in range(ag.shape[0]):
        tmp = 0
        for j in range(ab.shape[0]):
            tmp += match_arr(ag[i], ab[j])
        out[i] = tmp
    return out

@nb.njit(parallel=True)
def match_sets(Antigens, Antibodies):
    out = np.empty((Antigens.shape[0] * Antibodies.shape[0], Antigens.shape[1])) # 5000 x 40000
    # multiprocessing per antigen matrix, may want to move this as suits your data
    for i in nb.prange(Antigens.shape[0]):
        for j in range(Antibodies.shape[0]):
            out[j+(5*i)] = match_mat_sum(Antigens[i], Antibodies[j]) # need to figure out the index to avoid race conditions
    return out

我非常依赖 Numba。关键优化之一不是用np.equal() 检查整​​行的等价性，而是编写一个自定义函数match_arr()，一旦发现不匹配的元素就会中断。希望这能让我们跳过大量的比较。

时间对比：

%timeit match_arr(arr1, arr2)
314 ns ± 0.361 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

%timeit np.equal(arr1, arr2)
1.07 µs ± 5.35 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

match_mat_sum

此函数仅计算表示两个矩阵之间完全匹配的总和的中间步长（40,000 x 1 向量）。这一步减少了两个矩阵，例如：(m x n), (o x n) -> (m)

match_sets()

最后一个函数通过nb.prange 使用显式并行循环将此操作并行化。您可能希望根据数据的外观将此函数移动到不同的循环（例如，如果您有一个抗原矩阵，但有 5000 个抗体矩阵，您应该将 prange 移动到内部循环，否则您将不会利用并行化）。假数据假设一些抗原和一些抗体基质。

这里要注意的另一件重要事情是out 数组上的索引。为了避免竞争条件，每个显式循环都需要写入一个唯一的空间。同样，根据您的数据，您需要索引正确的“位置”来放置结果。

在具有 16 GB RAM 的 Ryzen 1600（6 核）上，使用这些假数据，我在 10.2 秒内生成了结果。

您的数据大约是 3200 倍。假设你有足够的内存，假设线性缩放，完整集大约需要 9 个小时。

您也可以编写某种批处理加载器，而不是将 5000 个巨型矩阵直接加载到内存中。

Answer 2

这个问题可以通过混合使用 numpy 广播和模块 numexpr 来解决，该模块可以快速执行操作，同时最小化中间值的存储

import numexpr as ne

# expand arrays dimensions to support broadcasting when doing comparison
Antigens, Antibodies = Antigens[None, :, :], Antibodies[:, None, :]
output = ne.evaluate('sum((Antigens==Antibodies)*1, axis=2)')
# *1 is a hack because numexpr does not currently support sum on bool

这可能比您当前的解决方案更快，但对于如此大的数组，这将需要一段时间。

numexpr 这个操作的性能有点差，但你至少可以在循环内使用广播：

output = np.zeros((Antibodies.shape[0],)*2, dtype=np.int32)
for row, out_row in zip(Antibodies, output):
    (row[None,:]==Antigens).sum(1, out=out_row)