Невозможно обучить VAE с деконволюционным слоем - PullRequest
Новая
       27

Невозможно обучить VAE с деконволюционным слоем

0 голосов
/ 03 сентября 2018

Я экспериментировал с реализацией VAE в Tensorflow для набора данных MNIST. Для начала я обучил VAE на основе кодера и декодера MLP. Он тренируется очень хорошо, потери уменьшаются, и он генерирует правдоподобно выглядящие цифры. Вот код декодера этого VAE на основе MLP: 

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, оптимизатор Адама, скорость обучения 1e-4.

EDIT:

Как указывал @Agost, мне, возможно, следует уточнить кое-что о моей функции потери и процессе обучения. Я моделирую апостериор как распределение Бернулли и максимизирую ELBO как свою потерю. Вдохновленный этой записью. Вот полный код кодера, декодера и потери: 

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)

1 Ответ

0 голосов
/ 03 сентября 2018

Может быть, вы используете sigmoid в последнем слое deconv, ограничивая вывод до 0-1, вы не делаете это в автоэнкодере на основе MLP или при добавлении полностью подключенного после deconvs столь возможной проблемы с диапазоном данных?

...