tensorflow - VQ-VAE-2论文的实现

Question

我正在尝试构建一个 2 阶段 VQ-VAE-2 + PixelCNN，如论文中所示:“使用 VQ-VAE-2 生成多样化的高保真图像”(https://arxiv.org/pdf/1906.00446.pdf)。我有 3 个实现问题:

1. 论文提到:

   We allow each level in the hierarchy to separately depend on pixels.

我理解 VQ-VAE-2 中的第二个潜在空间必须 以第一个潜在空间和一个的串联为条件 图像的下采样版本。这是正确的吗？

论文“Conditional Image Generation with PixelCNN Decoders”( https://papers.nips.cc/paper/6527-conditional-image-generation-with-pixelcnn-decoders.pdf ) 说道:

h is a one-hot encoding that specifies a class this is equivalent to adding a class dependent bias at every layer.

据我了解，条件作为一维张量输入，通过卷积注入(inject)偏差。现在，对于 2 阶段条件 PixelCNN，需要以类向量为条件，还需要以前一阶段的潜在代码为条件。我看到的一种可能性是附加它们并提供 3D 张量。有没有人看到另一种方法来做到这一点？

损失和优化在 2 个阶段中保持不变。只需将每个阶段的损失添加到优化后的最终损失中即可。这是正确的吗？

Answer 1

与该论文的一位作者讨论时，我收到了所有这些问题的答案并在下面分享。

问题1

这是正确的，但是图像的下采样是通过跨步卷积而不是非参数调整大小来实现的。这可以被吸收为编码器架构的一部分，如下所示(每个变量后面的数字表示它们的空间暗淡，例如 h64 是 [B, 64, 64, D] 等等)。

   h128 = Relu(Conv2D(image256, stride=(2, 2)))
   h64 = Relu(Conv2D(h128, stride=(2, 2)))
   h64 = ResNet(h64)

现在为了获取 h32 和 q32 我们可以这样做:

   h32 = Relu(Conv2D(h64,  stride=(2, 2)))
   h32 = ResNet(h32)
   q32 = Quantize(h32)

这样，渐变就会一直流回到图像，因此我们在 h32 和 image256 之间存在依赖关系。

在任何地方，您都可以使用 1x1 卷积来调整最后一个维度(特征层)的大小，使用跨步卷积进行下采样，使用跨步转置卷积进行空间维度上采样。因此，对于这个量化底层的示例，您需要首先在空间上对 q32 进行上采样，使其变为 64x64，然后将其与 h64 组合，并将结果提供给量化器。为了获得额外的表现力，我们还在中间插入了一个剩余堆栈。它看起来像这样:

    hq32 = ResNet(Conv2D(q32, (1, 1)))
    hq64 = Conv2DTranspose(hq32, stride=(2, 2))
    h64 = Conv2D(concat([h64, hq64]), (1, 1))
    q64 = Quantize(h64)

问题2

原始 PixelCNN 论文还描述了如何使用卷积来使用空间调节。将类嵌入扁平化和附加作为全局调节并不是一个好主意。您想要做的是应用转置卷积来对齐空间维度，然后应用 1x1 卷积来将特征维度与 Pixelcnn 的隐藏代表进行匹配，然后将其相加。

问题3

分开训练它们是个好主意。除了隔离损失等并能够为每个阶段调整适当的学习率之外，您还可以在每个阶段使用 GPU/TPU 的全部内存容量。这些先验随着规模的扩大做得越来越好，所以最好不要否认这一点。