Numpy скалярное произведение по заданным c осям

Question

У меня есть массив изображений размером 512x512, и я хочу выполнять операции с блоками 8x8. На данный момент у меня есть что-то вроде этого:

output = np.zeros(512, 512)

for i in range(0, 512, 8):
    for j in rangerange(0, 512, 8):
        a = input[i:i+8, j:j+8]
        b = some_other_array[i:i+8, j:j+8]
        output[i:i+8, j:j+8] = np.dot(a, b)

, где a & b - блоки 8x8, полученные из исходного массива. Я хотел бы ускорить этот код, используя векторизованные операции. Я изменил свои входные данные следующим образом:

input = input.reshape(64, 8, 64, 8)
some_other_array = some_other_array.reshape(64, 8, 64, 8)

Как я могу выполнить скалярное произведение только на осях 1 & 3, чтобы вывести массив формы (64, 8, 64, 8)?

Я пробовал np.tensordot(input, some_other_array, axes=([0, 1], [2, 3])), который дает правильную форму вывода, но значения не соответствуют выводу из l oop выше. Я также смотрел np.einsum, но я не встречал простого примера того, чего я пытаюсь достичь.

Answer 1

Как вы и подозревали, np.einsum может позаботиться об этом. Если input и some_other_array имеют формы (64, 8, 64, 8), тогда, если вы напишете

output = np.einsum('ijkl,ilkm->ijkm', input, some_other_array)  

, тогда output также будет иметь форму (64, 8, 64, 8), где умножение матриц (т.е. np.dot) было выполняется только по осям 1 и 3.

Строковый аргумент для np.einsum выглядит сложным, но на самом деле это комбинация двух вещей. Во-первых, матричное умножение задается как jl,lm->jm (см., Например, этот ответ на einsum ). Во-вторых, мы не хотим ничего делать с осями 0 и 2, поэтому для них я просто пишу ik,ik->ik. Объединение двух дает ijkl,ilkm->ijkm.

Answer 2

Они будут работать, если вы немного измените их порядок. Если input и some_other_array имеют форму (64,8,64,8), тогда:

input = input.transpose(0,2,1,3)
some_other_array = some_other_array.transpose(0,2,1,3)

Это изменит их порядок на 64,64,8,8. На этом этапе вы можете вычислить матричное умножение. Обратите внимание, что вам нужен matmul для вычисления блочных произведений, а не точка, которая попытается умножить все матрицы.

output = input @ some_other_array
output = output.transpose(0,2,1,3)
output = output.reshape(512,512)