Умножить массив матриц на массив векторов в Numpy
У меня есть массив матриц A формы A.shape = (N, 3, 3) и массив векторов V формы V.shape = (N, 3). Я хочу получить массив (N, 3), где каждый вектор является результатом умножения i-й матрицы на i-й вектор. Такой как:
result = [A[i] @ V[i] for i in range(N)]
Есть ли способ сделать это без зацикливания с помощью NumPy? В идеале мне бы также хотелось, чтобы это работало с массивами более высоких размерностей, например A.shape = (N, M, 3, 3) и V.shape = (N, M, 3, 3).
Я только что реорганизовал весь свой код, чтобы иметь такое расположение памяти из-за правил трансляции numpy. Насколько я понимаю, numpy всегда использует последние измерения для работы с трансляцией ((N, M, 3) + (3,) нормально, но (3, N, M) + (3,) выдает ошибку). Почему другой макет быстрее?
Потому что он более дружелюбен к SIMD. Модули SIMD на современных процессорах x86-64 обычно могут вычислять 16 32-битных чисел с плавающей запятой за цикл или 8 64-битных чисел с плавающей запятой. При нынешней компоновке можно использовать только часть модулей SIMD (всего 3 шт.). Что касается Numpy, он не оптимизирован для работы на оси с небольшим количеством элементов (и, конечно, никогда не будет).
Ответы 1
Я бы рекомендовал использовать np.einsum(). Это позволяет эффективно умножать матрицы по нескольким измерениям без явных циклов. (ссылка)
Вот как вы можете это использовать:
import numpy as np
def matrix_vector_multiply(A, V):
return np.einsum('...ij,...j->...i', A, V)
Это решение работает как для случая, когда A.shape = (N, 3, 3) и V.shape = (N, 3), так и для более высоких измерений, таких как A.shape = (N, M, 3, 3) и V.shape = (N, M, 3).
Я протестировал вышеуказанную функцию и сравнил ее с вашим зацикленным решением:
N = 5
A1 = np.random.rand(N, 3, 3)
V1 = np.random.rand(N, 3)
result1_loop = np.array([A1[i] @ V1[i] for i in range(N)])
result1_einsum = matrix_vector_multiply(A1, V1)
print("Loop result shape:", result1_loop.shape)
print("Einsum result shape:", result1_einsum.shape)
print("Results match:", np.allclose(result1_loop, result1_einsum))
Вот результат:
Форма результата цикла: (5, 3) Форма результата Einsum: (5, 3) Результаты совпадают: Истинный
Этот пост также поможет лучше понять einsum
Обратите внимание, что Numpy неэффективен для работы с матрицами 3x3. Нативный код или код Numba может работать намного быстрее, чем Numpy. Я думаю, что использование лучшего макета памяти, например
3x3xNxM, может сделать это еще быстрее (как с Numba, так и с Numpy).