Проблема состоит в том, что numpy рассматривает многомерные массивы как стеки матриц, и всегда последние два измерения предполагаются линейными измерениями пространства. Это означает, что точечный продукт не будет работать, если свернуть измерение first вашего 3d-массива.
Вместо этого самое простое, что вы можете сделать, это преобразовать ваш 3d-массив в 2d, выполнивумножение матрицы и преобразование обратно в трехмерный массив. Это также будет использовать оптимизированный код BLAS, который является одним из больших преимуществ numpy.
import numpy as np
S_pinv = np.random.rand(3, 4)
images = np.random.rand(4, 5, 6)
# error:
# (S_pinv @ images).shape
res_shape = S_pinv.shape[:1] + images.shape[1:] # (3, 5, 6)
res = (S_pinv @ images.reshape(images.shape[0], -1)).reshape(res_shape)
print(res.shape) # (3, 5, 6)
Так что вместо (3,n) x (n,h,w) мы делаем (3,n) x (n, h*w) -> (3, h*w), который мы преобразуем обратно в (3, h, w). Изменение формы является бесплатным, потому что это не означает каких-либо реальных манипуляций с данными в памяти (только переосмысление единственного блока памяти, лежащего в основе массива), и, как я сказал, правильные матричные продукты сильно оптимизированы в numpy.
Поскольку вы попросили более интуитивно понятный способ , вот альтернативный вариант использования numpy.einsum. Возможно, он будет медленнее, но он будет очень прозрачным, если вы немного привыкнете к его обозначениям:
res_einsum = np.einsum('tn,nhw -> thw', S_pinv, images)
print(np.array_equal(res, res_einsum)) # True
В этой записи обозначены все измерения входных массивов: для S_pinv первое и второеразмеры называются t и n, соответственно, и аналогично n, h и w для images. Выходные данные имеют размеры thw, что означает, что любые оставшиеся измерения, отсутствующие в выходной форме, будут суммироваться после умножения входных массивов. Это именно то, что вам нужно.
Как вы отметили в комментарии, вы также можете транспонировать свои массивы так, чтобы np.dot находил нужные измерения в нужном месте. Но это также будет медленным, потому что это может привести к копированию в памяти или, по крайней мере, к неоптимальному циклу над вашими массивами.
Я сделал быстрое сравнение по времени, используя следующие определения:
def reshaped(S_pinv, images):
res_shape = S_pinv.shape[:1] + images.shape[1:]
return (S_pinv @ images.reshape(images.shape[0], -1)).reshape(res_shape)
def einsummed(S_pinv, images):
return np.einsum('tn,nhw -> thw', S_pinv, images)
def transposed(S_pinv, images):
return (S_pinv @ images.transpose(2, 0, 1)).transpose(1, 2, 0)
А вот временной тест, использующий встроенную в IPython магию %timeit и некоторые более реалистичные размеры массивов:
>>> S_pinv = np.random.rand(3, 100)
... images = np.random.rand(100, 200, 300)
... args = S_pinv, images
... %timeit reshaped(*args)
... %timeit einsummed(*args)
... %timeit transposed(*args)
5.92 ms ± 460 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
15.9 ms ± 190 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
44.5 ms ± 329 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)