How to implement tf.space_to_depth with numpy?

Viewed 2458

These is a function in tensorflow called tf.space_to_depth. The implementation of this function in Tensorflow source code is so difficult for me. Could you please help me to implement it using numpy?

Here is some code to visualize how this function works. By the way, before everything, it is good to mention that the input to tensorflow's function should have input shape : [batch, height, width, depth]

Assuming this code. First we need to define a tensor:

norm = tf.reshape(tf.range(0,72),(1,6,6,2))

Here is the value of depth 1 (norm[0,:,:,0]):

[[ 0,  2,  4,  6,  8, 10],
 [12, 14, 16, 18, 20, 22],
 [24, 26, 28, 30, 32, 34],
 [36, 38, 40, 42, 44, 46],
 [48, 50, 52, 54, 56, 58],
 [60, 62, 64, 66, 68, 70]]

Here is the value of depth 2 (norm[0,:,:,1]):

[[ 1,  3,  5,  7,  9, 11],
 [13, 15, 17, 19, 21, 23],
 [25, 27, 29, 31, 33, 35],
 [37, 39, 41, 43, 45, 47],
 [49, 51, 53, 55, 57, 59],
 [61, 63, 65, 67, 69, 71]]

In the next step, I would like to apply tf.space_to_depth function and here it is:

trans = tf.space_to_depth(norm,2)

Output shape is : (1,3,3,8) and here is the output of this function:

trans[0,:,:,0]
[[ 0,  4,  8],
 [24, 28, 32],
 [48, 52, 56]]
trans[0,:,:,1]
[[ 1,  5,  9],
 [25, 29, 33],
 [49, 53, 57]]
trans[0,:,:,2]
[[ 2,  6, 10],
 [26, 30, 34],
 [50, 54, 58]]
trans[0,:,:,3]
[[ 3,  7, 11],
 [27, 31, 35],
 [51, 55, 59]]
trans[0,:,:,4]
[[12, 16, 20],
 [36, 40, 44],
 [60, 64, 68]]
trans[0,:,:,5]
[[13, 17, 21],
 [37, 41, 45],
 [61, 65, 69]]
trans[0,:,:,6]
[[14, 18, 22],
 [38, 42, 46],
 [62, 66, 70]]
trans[0,:,:,7]
[[15, 19, 23],
 [39, 43, 47],
 [63, 67, 71]]

May someone help me how can I implement a vectorized version of this function in numpy?

Appreciating in advance for any response!

1 Answers
Related