tf.space_to_depth() 如何在张量流中工作？

Question

我是 pytorch 用户。我在 tensorflow 中有一个预训练模型，我想将它转移到 pytorch 中。在模型架构的一部分中，我的意思是在 tensorflow 定义的模型中，有一个函数 tf.space_to_depth 将输入大小 (None, 38,38,64) 转换为 (None, 19, 19, 256)。 (https://www.tensorflow.org/api_docs/python/tf/space_to_depth) 是这个函数的文档。但我无法理解这个函数实际上是做什么的。能否请您提供一些 numpy 代码来为我说明？

其实我想在 pytorch 中做一个完全相似的层。

tensorflow 中的一些代码揭示了另一个秘密： 这是一些代码：

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)

这是输出：

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

正如你在上面看到的，除了数据重塑之外，张量值也发生了变化！

Answer 1

这个 tf.space_to_depth 将您的输入分成块并将它们连接起来。

在您的示例中，输入是 38x38x64（我猜 block_size 是 2）。因此，该函数将您的输入分成 4 (block_size x block_size) 并将它们连接起来，从而得到 19x19x256 的输出。

您只需将每个通道（输入）划分为 block_size*block_size 块（每个块的大小为 width/block_size x height/block_size）并连接所有这些块。使用 numpy 应该非常简单。

希望对你有帮助。

Answer 2

PyTorch 的一个很好的参考是 PixelShuffle 模块 here 的实现。这显示了与 Tensorflow 的depth_to_space 等效的实现。基于此，我们可以实现缩放因子小于 1 的 pixel_shuffle，这就像 space_to_depth。例如，downscale_factor=0.5 就像 space_to_depth 与 block_size=2。

def pixel_shuffle_down(input, downscale_factor):
    batch_size, channels, in_height, in_width = input.size()
    out_channels = channels / (downscale_factor ** 2)
    block_size = 1 / downscale_factor

    out_height = in_height * downscale_factor
    out_width = in_width * downscale_factor

    input_view = input.contiguous().view(
        batch_size, channels, out_height, block_size, out_width, block_size)

    shuffle_out = input_view.permute(0, 1, 3, 5, 2, 4).contiguous()
    return shuffle_out.view(batch_size, out_channels, out_height, out_width)

注意：我还没有验证这个实现，我不确定它是否与 pixel_shuffle 完全相反，但这是基本思想。我还在 PyTorch Github 上打开了一个关于 here 的问题。在 NumPy 中，等效代码将分别使用 reshape 和 transpose 而不是 view 和 permute。

Answer 3

结论：tf.space_to_depth() 仅输出输入张量的副本，其中将高度和宽度维度的值移动到深度维度。

如果你稍微修改你的代码，像这样

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)

那么你会得到以下结果：

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

希望对你有帮助。

Answer 4

在 Pytorch 中使用 split 和 stack 函数以及 permute 可以得到与 tensorflow 中的 space_to_depth 相同的结果。这是 Pytorch 中的代码。 假设输入是 BHWC 格式。

根据block_size和输入形状，我们可以计算出输出形状。 首先，它将“宽度”维度或维度 #2 上的输入拆分为 block_size。此操作的结果是一个长度为 d_width 的数组。这就像您将蛋糕（按块大小）切成 d_width 块。 然后对于每一块，你重塑它，使其具有正确的输出高度和输出深度（通道）。最后，我们将这些部分堆叠在一起并执行排列。

希望对您有所帮助。

def space_to_depth(input, block_size)
    block_size_sq = block_size*block_size
    (batch_size, s_height, s_width, s_depth) = input.size()
    d_depth = s_depth * self.block_size_sq
    d_width = int(s_width / self.block_size)
    d_height = int(s_height / self.block_size)
    t_1 = input.split(self.block_size, 2)
    stack = [t_t.contiguous().view(batch_size, d_height, d_depth) for t_t in t_1]
    output = torch.stack(stack, 1)
    output = output.permute(0, 2, 1, 3)
    return output

Answer 5

也许这个可行：

sudo apt install nvidia-cuda-toolkit

它对我有用。