VQ-VAE-2论文的实现

Question

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

1. 论文中提到：

   我们允许层次结构中的每个级别单独依赖于像素。

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

2. 论文“Conditional Image Generation with PixelCNN Decoders”（https://papers.nips.cc/paper/6527-conditional-image- Generation-with-pixelcnn-decoders.pdf）说：

   h 是一种 one-hot 编码，它指定一个类，这相当于在每一层添加一个类相关偏差。

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

3. 损失和优化在两个阶段都没有变化。只需将每个阶段的损失添加到优化后的最终损失中即可。那是对的吗 ？

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 的全部内存容量。这些先验随着规模的扩大做得越来越好，所以最好不要否认他们这一点。