how tf.space_to_depth() works in tensorflow?

Question

I am a pytorch user. I have got a pretrained model in tensorflow and I would like to transfer it into pytorch. In one part of model architecture, I mean in tensorflow-defined model, there is a function tf.space_to_depth which transfers an input size of (None, 38,38,64) to (None, 19,19, 256). (https://www.tensorflow.org/api_docs/python/tf/space_to_depth) is the doc of this function. But I could not understand what this function actually do. Could you please provide some numpy codes to illustrate it for me?

Actually I would like to make an exact similar layer in pytorch.

Some codes in tensorflow reveals another secret: Here is some codes:

import numpy as np
import tensorflow as tf

norm = tf.random_normal([1, 2, 2, 1], mean=0, stddev=1)
trans = tf.space_to_depth(norm,2)

with tf.Session() as s:
    norm = s.run(norm)
    trans = s.run(trans)



print("Norm")
print(norm.shape)
for index,value in np.ndenumerate(norm):
    print(value)

print("Trans")
print(trans.shape)
for index,value in np.ndenumerate(trans):
    print(value)

And here is the output:

Norm
(1, 2, 2, 1)
0.695261
0.455764
1.04699
-0.237587
Trans
(1, 1, 1, 4)
1.01139
0.898777
0.210135
2.36742

As you can see above, In Addition to data reshaping, the tensor values has changed!

Answer 1

This tf.space_to_depth divides your input into blocs and concatenates them.

In your example the input is 38x38x64 (and I guess the block_size is 2). So the function divides your input into 4 (block_size x block_size) and concatenates them which gives your 19x19x256 output.

You just need to divide each of your channel (input) into block_size*block_size patches (each patch has a size of width/block_size x height/block_size) and concatenate all of these patches. Should be pretty straightforward with numpy.

Hope it helps.

Answer 2

Conclusion: tf.space_to_depth() only outputs a copy of the input tensor where values from the height and width dimensions are moved to the depth dimension.

If you modify your code a little bit, like this

norm = tf.random_normal([1, 2, 2, 1], mean=0, stddev=1)

with tf.Session() as s:
    norm = s.run(norm)

trans = tf.space_to_depth(norm,2)

with tf.Session() as s:
    trans = s.run(trans)

Then you will have the following results:

Norm
(1, 2, 2, 1)
-0.130227
2.04587
-0.077691
-0.112031
Trans
(1, 1, 1, 4)
-0.130227
2.04587
-0.077691
-0.112031

Hope this can help you.