使用 numpy 计算 3D 矩阵乘法的有效方法

Question

How can I efficiently write and calculate this multiplication using numpy:

 for k in range(K):
    for i in range(SIZE):
       for j in range(SIZE):
          for i_b in range(B_SIZE):
             for j_b in range(B_SIZE):
                for k_b in range(k+1):
                   data[k, i * w + i_b, j * h + j_b] += arr1[k_b, i_b, j_b] * arr2[k_b, i, j]

For example:

SIZE, B_SIZE = 32, 8
arr1.shape -> (8, 8, 8)
arr2.shape -> (8, 32, 32)
data.shape -> (K, 256, 256)

Thank you.

Answer 1

You can use Numba for such kind of non-trivial case and rework the loops to use efficiently the CPU cache . Here is an example:

import numba as nb

@nb.njit
def compute(data, arr1, arr2):
    for k in range(K):
        for k_b in range(k+1):
            for i in range(SIZE):
                for j in range(SIZE):
                    tmp = arr2[k_b, i, j]
                    for i_b in range(B_SIZE):
                        for j_b in range(B_SIZE):
                            data[k, i * w + i_b, j * h + j_b] += arr1[k_b, i_b, j_b] * tmp

If you do this operation once, then you can pre-compile the Numba code by providing the types of the arrays. If K is big, then you can parallelize the code using @nb.njit(parallel=True) and use for k in nb.prange(K) rather than for k in range(K) . This should be several order of magnitude fater.