3

I have a 512x512 image array and I want to perform operations on 8x8 blocks. At the moment I have something like this:

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)

where a & b are 8x8 blocks derived from the original array. I would like to speed up this code by using vectorised operations. I have reshaped my inputs like this:

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

How could I perform a dot product on only axes 1 & 3 to output an array of shape (64, 8, 64, 8)?

I have tried np.tensordot(input, some_other_array, axes=([0, 1], [2, 3])) which gives the correct output shape, but the values do not match the output from the loop above. I've also looked at np.einsum but I haven't come across a simple example with what I'm trying to achieve.

Anjum Sayed
  • 872
  • 9
  • 20
  • 1
    BTW: `for i in range(0, 512, 8)` and you don't need do `i*8 : (i+1)*8` but `i:i+8`. It will be more readable. – furas May 30 '20 at 08:49
  • [here's neat way to split your image into blocks](https://stackoverflow.com/a/16858283/3595907) – DrBwts May 30 '20 at 10:29

2 Answers2

3

As you suspected, np.einsum can take care of this. If input and some_other_array have shapes (64, 8, 64, 8), then if you write

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

then output will also have shape (64, 8, 64, 8), where matrix multiplication (i.e. np.dot) has been done only on axes 1 and 3.

The string argument to np.einsum looks complicated, but really it's a combination of two things. First, matrix multiplication is given by jl,lm->jm (see e.g. this answer on einsum). Second, we don't want to do anything to axis 0 and 2, so for them I just write ik,ik->ik. Combining the two gives ijkl,ilkm->ijkm.

macroeconomist
  • 671
  • 3
  • 8
2

They'll work if you reorder them a bit. If input and some_other_array are both shaped (64,8,64,8), then:

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

This will reorder them to 64,64,8,8. At this point you can compute a matrix multiplication. Do note that you need matmul to compute the block products, and not dot, which will try to multiply the entire matrices.

output = input @ some_other_array
output = output.transpose(0,2,1,3)
output = output.reshape(512,512)
Mercury
  • 3,417
  • 1
  • 10
  • 35
  • 1
    Although the `np.einsum` answer was accepted, it's worth noting that for my application this method was faster. Thanks! – Anjum Sayed May 30 '20 at 15:06