Векторизация по строкам numpy

Question

Я пишу нейронный net с нуля, и мне нужно выполнить следующую операцию: Для каждого ряда матрицы dY возьмите с собой внешнее произведение того же ряда другой матрицы S (той же формы, что и Y), и умножьте эту строку dY на эту внешнюю матрицу (S [i ,:], S [i ,:]). Также умножьте dY * S поэлементно и добавьте это к нему.

Код ниже делает это, но он не векторизован. Можете ли вы помочь мне ускорить это?

out = dY.copy()
for i in range(dY.shape[0]):
    out[i, :] = dY[i, :] * S[i, :] - dY[i, :].dot(np.outer(S[i, :], S[i, :]))

Обновление: следующее берет (n, m) матрицу S и возвращает матрицу формы (n, m, m), где для каждой строки, мы берем внешний продукт с собой.

np.einsum("ab,ad->abd", S, S)

Обновление 2: окончательно решил его, используя два приложения np.einsum.

S_outer = np.einsum("ab,ad->abd", S, S)
return dY * S - np.einsum("ab,abc->ac", dY, S_outer)

Answer 1

Публикация решения, которое я нашел в качестве ответа.

Вы можете сделать это с помощью двух вызовов np.einsum ().

S_outer = np.einsum("ab,ad->abd", S, S)
return dY * S - np.einsum("ab,abc->ac", dY, S_outer)

Answer 2

Давайте разбить это на части. У вас есть

out[i, :] = dY[i, :] * S[i, :] - dY[i, :].dot(np.outer(S[i, :], S[i, :]))

Давайте запишем это как

p = dY[i, :] * S[i, :]
si = S[i, :]
q = dY[i, :].dot(si[:, None] * si)
out[i, :] = p - q

Очевидно, что p может быть учтено из l oop полностью до dY * S. Вы можете вычислить q, получив стопку внешних продуктов в форме (S.shape[0], S.shape[1], S.shape[1]) и применив @, он же np.matmul. Ключ в том, что matmul передает все, кроме двух последних измерений, вместе, в то время как dot эффективно берет внешний продукт. Это позволяет вам указать, что dY является стеком векторов, а не матрицей, просто введя новое измерение:

out = dY * S - dY[:, None, :] @ (S[:, :, None] * S[:, None, :])