田海立@ 2020-10-20

TensorFlow算子space_to_depth是depth_to_space的逆操作。本文用图文的方式来解释该算子运算的方式。

 

一、space_to_depth原型

space_to_depth是把space数据（width和height维）移到depth（Channel）维上，与depth_to_space刚好是反向的操作。对应到ML该操作是把width和height维上各取block_size都分给depth上。所以，对应有一个参数block_size，要求原Tensor的height和width都是block_size的整数倍。

这样，新的Tensor

Width是 input_width / block_size；Height是 input_height / block_size；Depth是input_depth * block_size * block_size

原型如下：

tf.nn.space_to_depth( input, block_size, data_format='NHWC', name=None )

因为这里严格区分了C与H/W各维度，如果数据格式不是NHWC时，需要指定data_format。

 

二、space_to_depth程序实现

以[1, 6, 4, 3] space_to_depth(block_size = 2)为例：

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

 

 

执行space_to_depth(block_size = 2)之后：

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

 

 

三、space_to_depth对数据的处理

space_to_depth操作对数据的处理就是：

选取[in_batch, block_size, block_size, in_depth]为一个Tensor；对Step#1里的Tensor做reshape操作改变为[in_batch, 1, 1, in_channel * block_size * block_size]；按照先Width方向再Height方向的顺序同Step#1一样选择Tensor；对Step#3里的Tensor同Step#3一样做reshape操作；把Step#4和Step#4里得到的新Tensor拼接起来。

最后的Tensor的shape也就是[in_batch, in_height / block_size, in_width / block_size, in_channel * (block_size*block_size)]

 

上述的处理过程，一张图展示就是这样：

 

总结

本文分析了tf.nn.space_to_depth对Tensor的处理。维度上做变换；数据上把height和width的block_size*block_size数据做reshape，也就实现了depth上的扩充，然后再整体拼接在一起。

 

【其他相关话题】

depth_to_space是这一过程的逆过程，参看《图解TensorFlow op：tf.nn.depth_to_space》。

data_format也可以指定为其他格式，比如NCHW，也请自行分析，这里不再赘述。

thl789 认证博客专家 系统分析师 Android 人工智能