如何用numpy實現tf.space_to_depth？

Question

這是一個功能tensorflow稱為tf.space_to_depth 。 在Tensorflow源代碼中實現此功能對我來說非常困難。 你能幫我用numpy來實現嗎？

下面是一些可視化此函數如何工作的代碼。 順便說一句，在所有事情之前，最好提一下，tensorflow功能的輸入應該有輸入形狀： [batch, height, width, depth]

假設這段代碼。 首先我們需要定義一個張量：

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

這是深度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]]

這是深度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]]

在下一步中，我想應用tf.space_to_depth函數，這里是：

trans = tf.space_to_depth(norm,2)

輸出形狀為：（1,3,3,8），這是此函數的輸出：

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]]

有人可以幫助我如何在numpy中實現這個函數的矢量化版本？

提前感謝任何回復！

Answer 1

您可以通過對reshape()和swapaxes()函數的適當調用來實現space_to_depth ：

import numpy as np

def space_to_depth(x, block_size):
    x = np.asarray(x)
    batch, height, width, depth = x.shape
    reduced_height = height // block_size
    reduced_width = width // block_size
    y = x.reshape(batch, reduced_height, block_size,
                         reduced_width, block_size, depth)
    z = np.swapaxes(y, 2, 3).reshape(batch, reduced_height, reduced_width, -1)
    return z

以下是tf.space_to_depth文檔中的示例 ：

In [328]: x = [[[[1], [2]],
     ...:       [[3], [4]]]]
     ...: 

In [329]: space_to_depth(x, 2)
Out[329]: array([[[[1, 2, 3, 4]]]])

In [330]: x = [[[[1, 2, 3], [4, 5, 6]],
     ...:       [[7, 8, 9], [10, 11, 12]]]]
     ...: 

In [331]: space_to_depth(x, 2)
Out[331]: array([[[[ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12]]]])

In [332]: x = [[[[1],   [2],  [5],  [6]],
     ...:       [[3],   [4],  [7],  [8]],
     ...:       [[9],  [10], [13],  [14]],
     ...:       [[11], [12], [15],  [16]]]]
     ...: 

In [333]: space_to_depth(x, 2)
Out[333]: 
array([[[[ 1,  2,  3,  4],
         [ 5,  6,  7,  8]],

        [[ 9, 10, 11, 12],
         [13, 14, 15, 16]]]])

這是你的例子：

In [334]: norm = np.arange(72).reshape(1, 6, 6, 2)

In [335]: trans = space_to_depth(norm, 2)

In [336]: trans[0, :, :, 0]
Out[336]: 
array([[ 0,  4,  8],
       [24, 28, 32],
       [48, 52, 56]])

In [337]: trans[0, :, :, 1]
Out[337]: 
array([[ 1,  5,  9],
       [25, 29, 33],
       [49, 53, 57]])

In [338]: trans[0, :, :, 7]
Out[338]: 
array([[15, 19, 23],
       [39, 43, 47],
       [63, 67, 71]])