Press "Enter" to skip to content

图解TensorFlow op:tf.nn.depth_to_space

田海立@CSDN 2020-10-20

 

TensorFlow算子depth_to_space的官方说明文档写得很简单,从中获取不到真正的执行方式,对于复杂的情况也就无从得知运算结果。本文用图文的方式来解释该算子运算的方式。

 

零、基础知识

 

阅读本文,你需要知道下列基本知识,可以从笔者之前的博文中了解:

 

数据的NCHW/NHWC排布方式

 

用到数据的NHWC摆放顺序

 

Tensor Reshape操作

 

用到Reshape操作,以及Reshape操作不改变数据的摆放原则。

 

一、depth_to_space原型

 

depth_to_space是把depth维的数据移到space上,与space_to_depth刚好是反向的操作。对应到ML系统里Tensor的NHWC/NCHW数据:depth对应Channel;space对应Height和Width。而且该操作是把depth上的数据分给Height以及Width上。所以,对应有一个参数block_size,要求 原Tensor的depth是block_size的block_size^2倍

 

这样,新的Tensor

Width是 input_width x block_size;
Height是 input_height x block_size;
Depth是input_depth / (block_size * block_size)

原型如下:

 

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

 

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

 

二、depth_to_space举例

 

怎幺执行的呢?官方举了三个很简单的例子,因为第一个更简单,已经包含在第二个里,这里列出两个。

 

1. [1,1,1,12] -> [1, 2, 2, 3]

 

程序实现如下:

 

 

直观上看,上面程序的[1,1,1,12] -> [1, 2, 2, 3]的depth_to_space操作如下:

 

 

2. [1, 2, 2, 4] -> [1,4,4,1]

 

程序实现如下:

 

 

直观上看,上面程序的[1,2,2,4] -> [1, 4, 4, 1]的depth_to_space操作如下:

 

 

3. 分析

 

depth_to_space执行时 维度的变化 很容易理解:Width是 input_width x block_size;Height是 input_height x block_size;Depth是input_depth / (block_size * block_size)。

 

但是数据怎幺移动的呢?

 

第一个例子是把源Tensor的一个Channel上的数据按照Channel -> Width -> Height的顺序在目标tensor上重新摆放。那为什幺不是别的顺序摆放呢?

 

第二个例子:

先把源Tensor上的一个Channel上的数据按照Width -> Height在目标Tensor顺序排放;
然后再在源Tensor的Width方向上选择下一个Channel,同样数据按照Width -> Height在目标Tensor顺序排放;
再源Tensor的Height方向上选择Channel,同样数据按照Width -> Height在目标Tensor顺序排放。

为什幺不是一个Channel上取数据,然后直接把[1, 2, 3, 4]就放在一行上?这到底是什幺规律?

 

如果两个例子结合起来,也就是Channel更深(操作后Channel仍大于1),原Width与Height也都大于1,结果怎样?

 

三、复杂情况[1, 2, 3, 12]

 

我们就看更复杂的情况:

 

 

该Tensor也就是这样的:

 

 

执行,之后:

 

 

四、depth_to_space对数据的处理

 

depth_to_space数据排布理解起来就是把按Channel处理,把Channel的数据重新排布。这里会牵涉:

原Channel的选择顺序:NHWC格式已经选择了Channel了,接下来当然就是Width方向了,然后是Height方向。
数据在新的目标Tensor上怎幺摆放的?无论NHWC还是NCHW格式,存储数据都是要以1-D存放。这里的数据已经是1-D的,也就是把这数据放到NHWC的空间上,而shape改变了。这不就是reshape的操作吗。

所以, depth_to_space操作对数据的处理 就是:

 

 

    1. 按照Width -> Height的顺序选择Channel;

 

    1. 对一个Channel里的数据,用reshape操作改变为[in_batch, block_size, block_size, in_channel / (block_size*block_size)]

 

    1. 对reshape之后的数据,按照Width -> Height的顺序再拼起来。

 

 

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

 

上述的处理过程,用图展示就是这样的:

 

 

总结

 

本文分析了tf.nn.depth_to_space对Tensor的处理。维度上做变换;数据上按Channel取数据,然后做reshape,再拼接在一起。

 

【其他相关话题】

 

space_to_depth是这一过程的逆过程,可自行分析。

 

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

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注