无法使用反卷积层训练 VAE

Question

我正在 MNIST 数据集的 Tensorflow 中尝试 VAE 实现。首先，我训练了一个基于 MLP 编码器和解码器的 VAE。它训练得很好，损失减少了，并且生成了看似合理的数字。下面是这个基于 MLP 的 VAE 的解码器代码：

x = sampled_z
x = tf.layers.dense(x, 200, tf.nn.relu)
x = tf.layers.dense(x, 200, tf.nn.relu)
x = tf.layers.dense(x, np.prod(data_shape))
img = tf.reshape(x, [-1] + data_shape)

下一步，我决定添加卷积层。仅更改编码器效果很好，但是当我在解码器中使用反卷积（而不是 fc 层）时，我根本没有得到任何训练。损失函数永远不会减少，输出总是黑色的。这是反卷积解码器的代码：

x = tf.layers.dense(sampled_z, 24, tf.nn.relu)
x = tf.layers.dense(x, 7 * 7 * 64, tf.nn.relu)
x = tf.reshape(x, [-1, 7, 7, 64])
x = tf.layers.conv2d_transpose(x, 64, 3, 2, 'SAME', activation=tf.nn.relu)
x = tf.layers.conv2d_transpose(x, 32, 3, 2, 'SAME', activation=tf.nn.relu)
x = tf.layers.conv2d_transpose(x, 1, 3, 1, 'SAME', activation=tf.nn.sigmoid)
img = tf.reshape(x, [-1, 28, 28])

这看起来很奇怪，代码对我来说看起来很好。我将其缩小到解码器中的反卷积层，那里有一些东西破坏了它。例如。如果我在最后一次反卷积之后添加一个全连接层（即使没有非线性！），它又可以工作了！代码如下：

x = tf.layers.dense(sampled_z, 24, tf.nn.relu)
x = tf.layers.dense(x, 7 * 7 * 64, tf.nn.relu)
x = tf.reshape(x, [-1, 7, 7, 64])
x = tf.layers.conv2d_transpose(x, 64, 3, 2, 'SAME', activation=tf.nn.relu)
x = tf.layers.conv2d_transpose(x, 32, 3, 2, 'SAME', activation=tf.nn.relu)
x = tf.layers.conv2d_transpose(x, 1, 3, 1, 'SAME', activation=tf.nn.sigmoid)
x = tf.contrib.layers.flatten(x)
x = tf.layers.dense(x, 28 * 28)
img = tf.reshape(x, [-1, 28, 28])

我真的有点卡在这一点上，有人知道这里会发生什么吗？我使用 tf 1.8.0，Adam 优化器，1e-4 学习率。

编辑：

正如@Agost 所指出的，我或许应该澄清一下我的损失函数和训练过程。我将后验建模为伯努利分布，并将 ELBO 最大化作为我的损失。灵感来自this 帖子。这是编码器、解码器和损失的完整代码：

def make_prior():
    mu = tf.zeros(N_LATENT)
    sigma = tf.ones(N_LATENT)
    return tf.contrib.distributions.MultivariateNormalDiag(mu, sigma)


def make_encoder(x_input):
    x_input = tf.reshape(x_input, shape=[-1, 28, 28, 1])
    x = conv(x_input, 32, 3, 2)
    x = conv(x, 64, 3, 2)
    x = conv(x, 128, 3, 2)
    x = tf.contrib.layers.flatten(x)
    mu = dense(x, N_LATENT)
    sigma = dense(x, N_LATENT, activation=tf.nn.softplus)  # softplus is log(exp(x) + 1)
    return tf.contrib.distributions.MultivariateNormalDiag(mu, sigma)    


def make_decoder(sampled_z):
    x = tf.layers.dense(sampled_z, 24, tf.nn.relu)
    x = tf.layers.dense(x, 7 * 7 * 64, tf.nn.relu)
    x = tf.reshape(x, [-1, 7, 7, 64])

    x = tf.layers.conv2d_transpose(x, 64, 3, 2, 'SAME', activation=tf.nn.relu)
    x = tf.layers.conv2d_transpose(x, 32, 3, 2, 'SAME', activation=tf.nn.relu)
    x = tf.layers.conv2d_transpose(x, 1, 3, 1, 'SAME')

    img = tf.reshape(x, [-1, 28, 28])

    img_distribution = tf.contrib.distributions.Bernoulli(img)
    img = img_distribution.probs
    img_distribution = tf.contrib.distributions.Independent(img_distribution, 2)
    return img, img_distribution


def main():
    mnist = input_data.read_data_sets(os.path.join(experiment_dir(EXPERIMENT), 'MNIST_data'))

    tf.reset_default_graph()

    batch_size = 128

    x_input = tf.placeholder(dtype=tf.float32, shape=[None, 28, 28], name='X')

    prior = make_prior()
    posterior = make_encoder(x_input)

    mu, sigma = posterior.mean(), posterior.stddev()

    z = posterior.sample()
    generated_img, output_distribution = make_decoder(z)

    likelihood = output_distribution.log_prob(x_input)
    divergence = tf.distributions.kl_divergence(posterior, prior)
    elbo = tf.reduce_mean(likelihood - divergence)
    loss = -elbo

    global_step = tf.train.get_or_create_global_step()
    optimizer = tf.train.AdamOptimizer(1e-3).minimize(loss, global_step=global_step)

Answer 1

可能是您在最终 deconv 层中使用 sigmoid 将输出限制为 0-1，您没有在基于 MLP 的自动编码器中执行此操作，或者在 deconvs 之后添加全连接时，可能存在数据范围问题？