在 python 中将 4D 数组与 2D 数组相乘和求和的最快方法？

Posted 2023-02-14

【中文标题】在 python 中将 4D 数组与 2D 数组相乘和求和的最快方法？

【英文标题】：Fastest way to multiply and sum 4D array with 2D array in python?

【发布时间】：2022-01-20 11:04:54

【问题描述】：

这是我的问题。我有两个矩阵A 和B，具有复杂的条目，维度分别为(n,n,m,m) 和(n,n)。

下面是我为得到一个矩阵C而执行的操作-

C = np.sum(B[:,:,None,None]*A, axis=(0,1))

计算一次以上大约需要 6-8 秒。因为我必须计算很多这样的Cs，所以需要很多时间。有没有更快的方法来做到这一点？ （我在多核 CPU 上使用 JAX NumPy 来做这些；普通的 NumPy 需要更长的时间）

n=77 和 m=512，如果您想知道的话。我可以在处理集群时进行并行化，但是数组的绝对大小会消耗大量内存。

【参考方案1】：

看起来你想要einsum:

C = np.einsum('ijkl,ij->kl', A, B)

在 Colab CPU 上使用 numpy 我得到了这个：

import numpy as np
x = np.random.rand(50, 50, 500, 500)
y = np.random.rand(50, 50)

def f1(x, y):
  return np.sum(y[:,:,None,None]*x, axis=(0,1))

def f2(x, y):
  return np.einsum('ijkl,ij->kl', x, y)

np.testing.assert_allclose(f1(x, y), f2(x, y))

%timeit f1(x, y)
# 1 loop, best of 5: 1.52 s per loop
%timeit f2(x, y)
# 1 loop, best of 5: 620 ms per loop