Extract corrections from the 3D matrix

Question

I have a three-dimensional matrix A size hxwxc . I want to extract ph x pw measurement patches from each "channel" c . ph divides h and pw divides w . In this example

 hxwxc = 4 x 4 x 3 ph x pw = 2 x 2 

I know how to do this in tenorflow using gather_nd , but I was hoping for something more efficient in terms of setting it up, because the sizes would be large, and I would prefer not to have an array of gather_nd indices in memory. Perhaps there is an intellectual adjustment? Any numpy or tensorflow solution would be very nice!

Answer 1

You can use some permutation and replacement of axes -

 A.reshape(h//ph,ph,w//pw,pw,-1).swapaxes(1,2) 

Run Example -

 In [46]: # Sample inputs ...: h,w,c = 10,12,3 ...: ph, pw = 2,2 ...: A = np.random.randint(0,9,(h,w,c)) ...: In [47]: A.reshape(h//ph,ph,w//pw,pw,-1).swapaxes(1,2).shape Out[47]: (5, 6, 2, 2, 3) 

Each element (as a block) along the first two axes represents patches. Thus. for the sample provided, we would have 5 x 6 = 30 corrections.

If you need these patches along one combined first axis, use another reshape -

 In [85]: out = A.reshape(h//ph,ph,w//pw,pw,-1).swapaxes(1,2).reshape(-1,ph,pw,c) In [86]: out.shape Out[86]: (30, 2, 2, 3) 

Let it be checked manually by checking the values ​​themselves -

 In [81]: A[:ph,:pw] # First patch Out[81]: array([[[6, 5, 2], [4, 0, 1]], [[0, 0, 4], [2, 3, 0]]]) In [82]: A[:ph,pw:2*pw] # Second patch Out[82]: array([[[8, 3, 3], [0, 0, 2]], [[8, 5, 4], [3, 4, 6]]]) In [83]: out[0] Out[83]: array([[[6, 5, 2], [4, 0, 1]], [[0, 0, 4], [2, 3, 0]]]) In [84]: out[1] Out[84]: array([[[8, 3, 3], [0, 0, 2]], [[8, 5, 4], [3, 4, 6]]])