Python: векторизация умножения матриц в циклах?

У меня есть массив N-by-M, в каждой записи которого мне нужно выполнить некоторые операции NumPy и поместить результат туда.

Прямо сейчас я делаю это наивным способом с двойной петлей:

import numpy as np

N = 10
M = 11
K = 100

result = np.zeros((N, M))

is_relevant = np.random.rand(N, M, K) > 0.5
weight = np.random.rand(3, 3, K)
values1 = np.random.rand(3, 3, K)
values2 = np.random.rand(3, 3, K)

for i in range(N):
    for j in range(M):
        selector = is_relevant[i, j, :]
        result[i, j] = np.sum(
            np.multiply(
                np.multiply(
                    values1[..., selector],
                    values2[..., selector]
                ), weight[..., selector]
            )
        )

Поскольку все циклические операции являются просто операциями NumPy, я думаю, что должен быть способ сделать это быстрее или без циклов.

Всего 2 ответа


Мы можем использовать комбинацию np.einsum и np.tensordot -

a = np.einsum('ijk,ijk,ijk->k',values1,values2,weight)
out = np.tensordot(a,is_relevant,axes=(0,2))

Кроме того, с одним einsum -

np.einsum('ijk,ijk,ijk,lmk->lm',values1,values2,weight,is_relevant)

А с np.dot и einsum -

is_relevant.dot(np.einsum('ijk,ijk,ijk->k',values1,values2,weight))

Кроме того, np.einsum с флагом optimize в np.einsum , установив для него значение True чтобы использовать BLAS.

Сроки -

In [146]: %%timeit
     ...: a = np.einsum('ijk,ijk,ijk->k',values1,values2,weight)
     ...: out = np.tensordot(a,is_relevant,axes=(0,2))
10000 loops, best of 3: 121 µs per loop

In [147]: %timeit np.einsum('ijk,ijk,ijk,lmk->lm',values1,values2,weight,is_relevant)
1000 loops, best of 3: 851 µs per loop

In [148]: %timeit np.einsum('ijk,ijk,ijk,lmk->lm',values1,values2,weight,is_relevant,optimize=True)
1000 loops, best of 3: 347 µs per loop

In [156]: %timeit is_relevant.dot(np.einsum('ijk,ijk,ijk->k',values1,values2,weight))
10000 loops, best of 3: 58.6 µs per loop

Очень большие массивы

Для очень больших массивов мы можем использовать numexpr чтобы использовать multi-cores процессоры -

import numexpr as ne

a = np.einsum('ijk,ijk,ijk->k',values1,values2,weight)
out = np.empty((N, M))
for i in range(N):
    for j in range(M):
        out[i,j] = ne.evaluate('sum(is_relevant_ij*a)',{'is_relevant_ij':is_relevant[i,j], 'a':a})


Есть идеи?

10000