def train_model():

    # Setup session
    sess = tu.setup_training_session()

    ##########
    # Innputs
    ##########

    # Setup async input queue of real images
    X_real = du.read_celebA()

    # Noise
    batch_size = tf.shape(X_real)[0]
    z_noise_for_D = tf.random_uniform((batch_size, FLAGS.z_dim,), minval=-1, maxval=1, name="z_input_D")
    z_noise_for_G = tf.random_uniform((batch_size, FLAGS.z_dim,), minval=-1, maxval=1, name="z_input_G")

    # k factor
    k_factor = tf.Variable(initial_value=0., trainable=False, name='anneal_factor')

    # learning rate
    lr = tf.Variable(initial_value=FLAGS.learning_rate, trainable=False, name='learning_rate')

    ########################
    # Instantiate models
    ########################
    G = models.Generator(nb_filters=FLAGS.nb_filters_G)
    D = models.Discriminator(h_dim=FLAGS.h_dim, nb_filters=FLAGS.nb_filters_D)

    ##########
    # Outputs
    ##########
    X_rec_real = D(X_real, output_name="X_rec_real")

    X_fake_for_D = G(z_noise_for_D, output_name="X_fake_for_D")
    X_rec_fake_for_D = D(X_fake_for_D, reuse=True, output_name="X_rec_fake_for_D")

    X_fake_for_G = G(z_noise_for_G, reuse=True, output_name="X_fake_for_G")
    X_rec_fake_for_G = D(X_fake_for_G, reuse=True, output_name="X_rec_fake_for_G")

    # output images for plots
    real_toplot = du.unnormalize_image(X_real, name="real_toplot")
    generated_toplot = du.unnormalize_image(X_fake_for_G, name="generated_toplot")
    real_rec_toplot = du.unnormalize_image(X_rec_real, name="rec_toplot")
    generated_rec_toplot = du.unnormalize_image(X_rec_fake_for_G, name="generated_rec_toplot")

    ###########################
    # Instantiate optimizers
    ###########################
    opt = tf.train.AdamOptimizer(learning_rate=lr, name='opt')

    ###########################
    # losses
    ###########################

    loss_real = losses.mae(X_real, X_rec_real)
    loss_fake_for_D = losses.mae(X_fake_for_D, X_rec_fake_for_D)
    loss_fake_for_G = losses.mae(X_fake_for_G, X_rec_fake_for_G)

    L_D = loss_real - k_factor * loss_fake_for_D
    L_G = loss_fake_for_G
    Convergence = loss_real + tf.abs(FLAGS.gamma * loss_real - loss_fake_for_G)

    ###########################
    # Compute updates ops
    ###########################

    dict_G_vars = G.get_trainable_variables()
    G_vars = [dict_G_vars[k] for k in dict_G_vars.keys()]

    dict_D_vars = D.get_trainable_variables()
    D_vars = [dict_D_vars[k] for k in dict_D_vars.keys()]

    G_gradvar = opt.compute_gradients(L_G, var_list=G_vars)
    G_update = opt.apply_gradients(G_gradvar, name='G_loss_minimize')

    D_gradvar = opt.compute_gradients(L_D, var_list=D_vars)
    D_update = opt.apply_gradients(D_gradvar, name='D_loss_minimize')

    update_k_factor = tf.assign(k_factor, k_factor + FLAGS.lambdak * (FLAGS.gamma * loss_real - loss_fake_for_G))
    update_lr = tf.assign(lr, tf.maximum(1E-6, lr / 2))

    ##########################
    # Summary ops
    ##########################

    # Add summary for gradients
    tu.add_gradient_summary(G_gradvar)
    tu.add_gradient_summary(D_gradvar)

    # Add scalar symmaries for G
    tf.summary.scalar("G loss", L_G)
    # Add scalar symmaries for D
    tf.summary.scalar("D loss", L_D)
    # Add scalar symmaries for D
    tf.summary.scalar("k_factor", k_factor)
    tf.summary.scalar("Convergence", Convergence)
    tf.summary.scalar("learning rate", lr)

    summary_op = tf.summary.merge_all()

    ############################
    # Start training
    ############################

    # Initialize session
    saver = tu.initialize_session(sess)

    # Start queues
    coord, threads = du.manage_queues(sess)

    # Summaries
    writer = tu.manage_summaries(sess)

    # Run checks on data dimensions
    list_data = [z_noise_for_D, z_noise_for_G]
    list_data += [X_real, X_rec_real, X_fake_for_G, X_rec_fake_for_G, X_fake_for_D, X_rec_fake_for_D]
    list_data += [generated_toplot, real_toplot]
    output = sess.run(list_data)
    tu.check_data(output, list_data)

    for e in tqdm(range(FLAGS.nb_epoch), desc="Training progress"):

        # Anneal learning rate
        if (e + 1) % 200 == 0:
            sess.run([update_lr])

        t = tqdm(range(FLAGS.nb_batch_per_epoch), desc="Epoch %i" % e, mininterval=0.5)
        for batch_counter in t:

            output = sess.run([G_update, D_update, update_k_factor])

            if batch_counter % (FLAGS.nb_batch_per_epoch // (int(0.5 * FLAGS.nb_batch_per_epoch))) == 0:
                output = sess.run([summary_op])
                writer.add_summary(output[-1], e * FLAGS.nb_batch_per_epoch + batch_counter)

            t.set_description('Epoch %s:' % e)

        # Plot some generated images
        Xf, Xr, Xrrec, Xfrec = sess.run([generated_toplot, real_toplot, real_rec_toplot, generated_rec_toplot])
        vu.save_image(Xf, Xr, title="current_batch", e=e)
        vu.save_image(Xrrec, Xfrec, title="reconstruction", e=e)

        # Save session
        saver.save(sess, os.path.join(FLAGS.model_dir, "model"), global_step=e)

        # Show data statistics
        output = sess.run(list_data)
        tu.check_data(output, list_data)

    # Stop threads
    coord.request_stop()
    coord.join(threads)

    print('Finished training!')
def train_model():

    # Setup session
    sess = tu.setup_session()

    # Placeholder for data and Mnist iterator
    mnist = input_data.read_data_sets(FLAGS.raw_dir, one_hot=True)
    # images = tf.constant(mnist.train.images)
    if FLAGS.data_format == "NHWC":
        X_real = tf.placeholder(tf.float32, shape=[FLAGS.batch_size, 28, 28, 1])
    else:
        X_real = tf.placeholder(tf.float32, shape=[FLAGS.batch_size, 1, 28, 28])

    with tf.device('/cpu:0'):
        imgs = mnist.train.images.astype(np.float32)
        npts = imgs.shape[0]
        if FLAGS.data_format == "NHWC":
            imgs = imgs.reshape((npts, 28, 28, 1))
        else:
            imgs = imgs.reshape((npts, 1, 28, 28))
        imgs = (imgs - 0.5) / 0.5
        # input_images = tf.constant(imgs)

    #     image = tf.train.slice_input_producer([input_images], num_epochs=FLAGS.nb_epoch)
    #     X_real = tf.train.batch(image, batch_size=FLAGS.batch_size, num_threads=8)

    #######################
    # Instantiate generator
    #######################
    list_filters = [256, 1]
    list_strides = [2] * len(list_filters)
    list_kernel_size = [3] * len(list_filters)
    list_padding = ["SAME"] * len(list_filters)
    output_shape = X_real.get_shape().as_list()[1:]
    G = models.Generator(list_filters, list_kernel_size, list_strides, list_padding, output_shape,
                         batch_size=FLAGS.batch_size, dset="mnist", data_format=FLAGS.data_format)

    ###########################
    # Instantiate discriminator
    ###########################
    list_filters = [32, 64]
    list_strides = [2] * len(list_filters)
    list_kernel_size = [3] * len(list_filters)
    list_padding = ["SAME"] * len(list_filters)
    D = models.Discriminator(list_filters, list_kernel_size, list_strides, list_padding,
                             FLAGS.batch_size, data_format=FLAGS.data_format)

    ###########################
    # Instantiate optimizers
    ###########################
    G_opt = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate, name='G_opt', beta1=0.5)
    D_opt = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate, name='D_opt', beta1=0.5)

    ###########################
    # Instantiate model outputs
    ###########################

    # noise_input = tf.random_normal((FLAGS.batch_size, FLAGS.noise_dim,), stddev=0.1)
    noise_input = tf.random_uniform((FLAGS.batch_size, FLAGS.noise_dim,), minval=-1, maxval=1)
    X_fake = G(noise_input)

    # output images
    X_G_output = du.unnormalize_image(X_fake)
    X_real_output = du.unnormalize_image(X_real)

    D_real = D(X_real)
    D_fake = D(X_fake, reuse=True)

    ###########################
    # Instantiate losses
    ###########################

    G_loss = objectives.binary_cross_entropy_with_logits(D_fake, tf.ones_like(D_fake))
    D_loss_real = objectives.binary_cross_entropy_with_logits(D_real, tf.ones_like(D_real))
    D_loss_fake = objectives.binary_cross_entropy_with_logits(D_fake, tf.zeros_like(D_fake))

    D_loss = D_loss_real + D_loss_fake

    # ######################################################################
    # # Some parameters need to be updated (e.g. BN moving average/variance)
    # ######################################################################
    # from tensorflow.python.ops import control_flow_ops
    # update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    # with tf.control_dependencies(update_ops):
    #     barrier = tf.no_op(name='update_barrier')
    # D_loss = control_flow_ops.with_dependencies([barrier], D_loss)
    # G_loss = control_flow_ops.with_dependencies([barrier], G_loss)

    ###########################
    # Compute gradient updates
    ###########################

    dict_G_vars = G.get_trainable_variables()
    G_vars = [dict_G_vars[k] for k in dict_G_vars.keys()]

    dict_D_vars = D.get_trainable_variables()
    D_vars = [dict_D_vars[k] for k in dict_D_vars.keys()]

    G_gradvar = G_opt.compute_gradients(G_loss, var_list=G_vars)
    G_update = G_opt.apply_gradients(G_gradvar, name='G_loss_minimize')

    D_gradvar = D_opt.compute_gradients(D_loss, var_list=D_vars)
    D_update = D_opt.apply_gradients(D_gradvar, name='D_loss_minimize')

    ##########################
    # Group training ops
    ##########################
    train_ops = [G_update, D_update]
    loss_ops = [G_loss, D_loss, D_loss_real, D_loss_fake]

    ##########################
    # Summary ops
    ##########################

    # Add summary for gradients
    tu.add_gradient_summary(G_gradvar)
    tu.add_gradient_summary(D_gradvar)

    # Add scalar symmaries
    tf.summary.scalar("G loss", G_loss)
    tf.summary.scalar("D loss real", D_loss_real)
    tf.summary.scalar("D loss fake", D_loss_fake)

    summary_op = tf.summary.merge_all()

    ############################
    # Start training
    ############################

    # Initialize session
    saver = tu.initialize_session(sess)

    # Start queues
    coord = tu.manage_queues(sess)

    # Summaries
    writer = tu.manage_summaries(sess)

    for e in tqdm(range(FLAGS.nb_epoch), desc="\nTraining progress"):

        t = tqdm(range(FLAGS.nb_batch_per_epoch), desc="Epoch %i" % e, mininterval=0.5)
        for batch_counter in t:

            # X_batch, _ = mnist.train.next_batch(FLAGS.batch_size)
            # if FLAGS.data_format == "NHWC":
            #     X_batch = np.reshape(X_batch, [-1, 28, 28, 1])
            # else:
            #     X_batch = np.reshape(X_batch, [-1, 1, 28, 28])
            # X_batch = (X_batch - 0.5) / 0.5

            X_batch = du.sample_batch(imgs, FLAGS.batch_size)
            output = sess.run(train_ops + loss_ops + [summary_op], feed_dict={X_real: X_batch})

            if batch_counter % (FLAGS.nb_batch_per_epoch // 20) == 0:
                writer.add_summary(output[-1], e * FLAGS.nb_batch_per_epoch + batch_counter)
            lossG, lossDreal, lossDfake = [output[2], output[4], output[5]]

            t.set_description('Epoch %i: - G loss: %.2f D loss real: %.2f Dloss fake: %.2f' %
                              (e, lossG, lossDreal, lossDfake))

            # variables = [v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)]
            # bmean = [v for v in variables if v.name == "generator/conv2D_1_1/BatchNorm/moving_mean:0"]
            # print sess.run(bmean)
            # raw_input()

        # Plot some generated images
        # output = sess.run([X_G_output, X_real_output])
        output = sess.run([X_G_output, X_real_output], feed_dict={X_real: X_batch})
        vu.save_image(output, FLAGS.data_format, e)

        # Save session
        saver.save(sess, os.path.join(FLAGS.model_dir, "model"), global_step=e)

    print('Finished training!')
Esempio n. 3
0
def train_model():

    # Setup session
    sess = tu.setup_training_session()

    ##########
    # Innputs
    ##########

    # Setup async input queue of real images
    X_real = du.read_celebA()

    # Noise
    batch_size = tf.shape(X_real)[0]
    z_noise_for_D = tf.random_uniform((
        batch_size,
        FLAGS.z_dim,
    ),
                                      minval=-1,
                                      maxval=1,
                                      name="z_input_D")
    z_noise_for_G = tf.random_uniform((
        batch_size,
        FLAGS.z_dim,
    ),
                                      minval=-1,
                                      maxval=1,
                                      name="z_input_G")

    # k factor
    k_factor = tf.Variable(initial_value=0.,
                           trainable=False,
                           name='anneal_factor')

    # learning rate
    lr = tf.Variable(initial_value=FLAGS.learning_rate,
                     trainable=False,
                     name='learning_rate')

    ########################
    # Instantiate models
    ########################
    G = models.Generator(nb_filters=FLAGS.nb_filters_G)
    D = models.Discriminator(h_dim=FLAGS.h_dim, nb_filters=FLAGS.nb_filters_D)

    ##########
    # Outputs
    ##########
    X_rec_real = D(X_real, output_name="X_rec_real")

    X_fake_for_D = G(z_noise_for_D, output_name="X_fake_for_D")
    X_rec_fake_for_D = D(X_fake_for_D,
                         reuse=True,
                         output_name="X_rec_fake_for_D")

    X_fake_for_G = G(z_noise_for_G, reuse=True, output_name="X_fake_for_G")
    X_rec_fake_for_G = D(X_fake_for_G,
                         reuse=True,
                         output_name="X_rec_fake_for_G")

    # output images for plots
    real_toplot = du.unnormalize_image(X_real, name="real_toplot")
    generated_toplot = du.unnormalize_image(X_fake_for_G,
                                            name="generated_toplot")
    real_rec_toplot = du.unnormalize_image(X_rec_real, name="rec_toplot")
    generated_rec_toplot = du.unnormalize_image(X_rec_fake_for_G,
                                                name="generated_rec_toplot")

    ###########################
    # Instantiate optimizers
    ###########################
    opt = tf.train.AdamOptimizer(learning_rate=lr, name='opt')

    ###########################
    # losses
    ###########################

    loss_real = losses.mae(X_real, X_rec_real)
    loss_fake_for_D = losses.mae(X_fake_for_D, X_rec_fake_for_D)
    loss_fake_for_G = losses.mae(X_fake_for_G, X_rec_fake_for_G)

    L_D = loss_real - k_factor * loss_fake_for_D
    L_G = loss_fake_for_G
    Convergence = loss_real + tf.abs(FLAGS.gamma * loss_real - loss_fake_for_G)

    ###########################
    # Compute updates ops
    ###########################

    dict_G_vars = G.get_trainable_variables()
    G_vars = [dict_G_vars[k] for k in dict_G_vars.keys()]

    dict_D_vars = D.get_trainable_variables()
    D_vars = [dict_D_vars[k] for k in dict_D_vars.keys()]

    G_gradvar = opt.compute_gradients(L_G, var_list=G_vars)
    G_update = opt.apply_gradients(G_gradvar, name='G_loss_minimize')

    D_gradvar = opt.compute_gradients(L_D, var_list=D_vars)
    D_update = opt.apply_gradients(D_gradvar, name='D_loss_minimize')

    update_k_factor = tf.assign(
        k_factor,
        k_factor + FLAGS.lambdak * (FLAGS.gamma * loss_real - loss_fake_for_G))
    update_lr = tf.assign(lr, lr / 2)

    ##########################
    # Summary ops
    ##########################

    # Add summary for gradients
    tu.add_gradient_summary(G_gradvar)
    tu.add_gradient_summary(D_gradvar)

    # Add scalar symmaries for G
    tf.summary.scalar("G loss", L_G)
    # Add scalar symmaries for D
    tf.summary.scalar("D loss", L_D)
    # Add scalar symmaries for D
    tf.summary.scalar("k_factor", k_factor)
    tf.summary.scalar("Convergence", Convergence)
    tf.summary.scalar("learning rate", lr)

    summary_op = tf.summary.merge_all()

    ############################
    # Start training
    ############################

    # Initialize session
    saver = tu.initialize_session(sess)

    # Start queues
    coord, threads = du.manage_queues(sess)

    # Summaries
    writer = tu.manage_summaries(sess)

    # Run checks on data dimensions
    list_data = [z_noise_for_D, z_noise_for_G]
    list_data += [
        X_real, X_rec_real, X_fake_for_G, X_rec_fake_for_G, X_fake_for_D,
        X_rec_fake_for_D
    ]
    list_data += [generated_toplot, real_toplot]
    output = sess.run(list_data)
    tu.check_data(output, list_data)

    for e in tqdm(range(FLAGS.nb_epoch), desc="Training progress"):

        # Anneal learning rate every 5 epoch
        if (e + 1) % 5 == 0:
            sess.run([update_lr])

        t = tqdm(range(FLAGS.nb_batch_per_epoch),
                 desc="Epoch %i" % e,
                 mininterval=0.5)
        for batch_counter in t:

            output = sess.run([G_update, D_update, update_k_factor])

            if batch_counter % (FLAGS.nb_batch_per_epoch //
                                (int(0.5 * FLAGS.nb_batch_per_epoch))) == 0:
                output = sess.run([summary_op])
                writer.add_summary(
                    output[-1], e * FLAGS.nb_batch_per_epoch + batch_counter)

            t.set_description('Epoch %s:' % e)

        # Plot some generated images
        Xf, Xr, Xrrec, Xfrec = sess.run([
            generated_toplot, real_toplot, real_rec_toplot,
            generated_rec_toplot
        ])
        vu.save_image(Xf, Xr, title="current_batch", e=e)
        vu.save_image(Xrrec, Xfrec, title="reconstruction", e=e)

        # Save session
        saver.save(sess, os.path.join(FLAGS.model_dir, "model"), global_step=e)

        # Show data statistics
        output = sess.run(list_data)
        tu.check_data(output, list_data)

    # Stop threads
    coord.request_stop()
    coord.join(threads)

    print('Finished training!')
Esempio n. 4
0
def train_model():

    # Setup session
    sess = tu.setup_training_session()

    # Setup async input queue of real images
    X_real16, X_real32, X_real64 = du.read_celebA()

    #######################
    # Instantiate generators
    #######################
    G16 = models.G16()
    G32 = models.G32()
    G64 = models.G64()

    ###########################
    # Instantiate discriminators
    ###########################
    D16 = models.D16()
    D32 = models.D32()
    D64 = models.D64()

    ###########################
    # Instantiate optimizers
    ###########################
    G_opt = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate,
                                   name='G_opt',
                                   beta1=0.5)
    D_opt = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate,
                                   name='D_opt',
                                   beta1=0.5)

    ###########################
    # Instantiate model outputs
    ###########################

    # noise_input = tf.random_normal((FLAGS.batch_size, FLAGS.noise_dim,), stddev=0.1)
    noise_input = tf.random_uniform((
        FLAGS.batch_size,
        FLAGS.noise_dim,
    ),
                                    minval=-1,
                                    maxval=1)

    X_fake16 = G16(noise_input)

    D16_real = D16(X_real16, mode="D")
    X_feat16, D16_fake = D16(X_fake16, reuse=True, mode="G")

    X_fake32 = G32(X_fake16, X_feat16)

    D32_real = D32(X_real32, mode="D")
    X_feat32, D32_fake = D32(X_fake32, reuse=True, mode="G")

    X_fake64 = G64(X_fake32, X_feat32)

    D64_real = D64(X_real64)
    D64_fake = D64(X_fake64, reuse=True)

    # output images
    X_fake16_output = du.unnormalize_image(X_fake16)
    X_real16_output = du.unnormalize_image(X_real16)

    X_fake32_output = du.unnormalize_image(X_fake32)
    X_real32_output = du.unnormalize_image(X_real32)

    X_fake64_output = du.unnormalize_image(X_fake64)
    X_real64_output = du.unnormalize_image(X_real64)

    ###########################
    # Instantiate losses
    ###########################

    G16_loss = objectives.binary_cross_entropy_with_logits(
        D16_fake, tf.ones_like(D16_fake))
    G32_loss = objectives.binary_cross_entropy_with_logits(
        D32_fake, tf.ones_like(D32_fake))
    G64_loss = objectives.binary_cross_entropy_with_logits(
        D64_fake, tf.ones_like(D64_fake))
    G_loss = G16_loss + G32_loss + G64_loss

    # Fake losses
    D16_loss_fake = objectives.binary_cross_entropy_with_logits(
        D16_fake, tf.zeros_like(D16_fake))
    D32_loss_fake = objectives.binary_cross_entropy_with_logits(
        D32_fake, tf.zeros_like(D32_fake))
    D64_loss_fake = objectives.binary_cross_entropy_with_logits(
        D64_fake, tf.zeros_like(D64_fake))

    # Real losses
    D16_loss_real = objectives.binary_cross_entropy_with_logits(
        D16_real, tf.ones_like(D16_real))
    D32_loss_real = objectives.binary_cross_entropy_with_logits(
        D32_real, tf.ones_like(D32_real))
    D64_loss_real = objectives.binary_cross_entropy_with_logits(
        D64_real, tf.ones_like(D64_real))

    D_loss = D16_loss_real + D32_loss_real + D64_loss_real
    D_loss += D16_loss_fake + D32_loss_fake + D64_loss_fake

    ###########################
    # Compute gradient updates
    ###########################

    dict_G16_vars = G16.get_trainable_variables()
    G16_vars = [dict_G16_vars[k] for k in dict_G16_vars.keys()]
    dict_G32_vars = G32.get_trainable_variables()
    G32_vars = [dict_G32_vars[k] for k in dict_G32_vars.keys()]
    dict_G64_vars = G64.get_trainable_variables()
    G64_vars = [dict_G64_vars[k] for k in dict_G64_vars.keys()]
    G_vars = G16_vars + G32_vars + G64_vars

    dict_D16_vars = D16.get_trainable_variables()
    D16_vars = [dict_D16_vars[k] for k in dict_D16_vars.keys()]
    dict_D32_vars = D32.get_trainable_variables()
    D32_vars = [dict_D32_vars[k] for k in dict_D32_vars.keys()]
    dict_D64_vars = D64.get_trainable_variables()
    D64_vars = [dict_D64_vars[k] for k in dict_D64_vars.keys()]
    D_vars = D16_vars + D32_vars + D64_vars

    G_gradvar = G_opt.compute_gradients(G_loss,
                                        var_list=G_vars,
                                        colocate_gradients_with_ops=True)
    G_update = G_opt.apply_gradients(G_gradvar, name='G_loss_minimize')

    D_gradvar = D_opt.compute_gradients(D_loss,
                                        var_list=D_vars,
                                        colocate_gradients_with_ops=True)
    D_update = D_opt.apply_gradients(D_gradvar, name='D_loss_minimize')

    ##########################
    # Summary ops
    ##########################

    # Add summary for gradients
    tu.add_gradient_summary(G_gradvar)
    tu.add_gradient_summary(D_gradvar)

    # Add scalar symmaries
    tf.summary.scalar("G16 loss", G16_loss)
    tf.summary.scalar("G32 loss", G32_loss)
    tf.summary.scalar("G64 loss", G64_loss)

    # Real losses
    tf.summary.scalar("D16 loss real", D16_loss_real)
    tf.summary.scalar("D32 loss real", D32_loss_real)
    tf.summary.scalar("D64 loss real", D64_loss_real)

    # Fake losses
    tf.summary.scalar("D16 loss fake", D16_loss_fake)
    tf.summary.scalar("D32 loss fake", D32_loss_fake)
    tf.summary.scalar("D64 loss fake", D64_loss_fake)

    summary_op = tf.summary.merge_all()

    ############################
    # Start training
    ############################

    # Initialize session
    saver = tu.initialize_session(sess)

    # Start queues
    du.manage_queues(sess)

    # Summaries
    writer = tu.manage_summaries(sess)

    # Run checks on data dimensions
    list_data = [noise_input]
    list_data += [X_fake16, X_fake32, X_fake64]
    list_data += [X_fake16_output, X_fake32_output, X_fake64_output]
    output = sess.run(list_data)
    tu.check_data(output, list_data)

    for e in tqdm(range(FLAGS.nb_epoch), desc="Training progress"):

        t = tqdm(range(FLAGS.nb_batch_per_epoch),
                 desc="Epoch %i" % e,
                 mininterval=0.5)
        for batch_counter in t:

            # Update D
            output = sess.run([D_update])

            # Update G
            output = sess.run([G_update])

            if batch_counter % (FLAGS.nb_batch_per_epoch // 20) == 0:
                output = sess.run([summary_op])
                writer.add_summary(
                    output[-1], e * FLAGS.nb_batch_per_epoch + batch_counter)

            t.set_description('Epoch %i:' % e)

        # Plot some generated images
        output = sess.run([
            X_fake16_output,
            X_real16_output,
            X_fake32_output,
            X_real32_output,
            X_fake64_output,
            X_real64_output,
        ])
        vu.save_image(output[:2], e=e, title="size_16")
        vu.save_image(output[2:4], e=e, title="size_32")
        vu.save_image(output[4:6], e=e, title="size_64")

        # Save session
        saver.save(sess, os.path.join(FLAGS.model_dir, "model"), global_step=e)

        # Show data statistics
        output = sess.run(list_data)
        tu.check_data(output, list_data)

    print('Finished training!')
Esempio n. 5
0
def train_model():

    # Setup session
    sess = tu.setup_session()

    # Setup async input queue of real images
    #X_input = du.input_data(sess)
    #X_real, X_fake_in = X_input[0], X_input[1]
    X_real, X_fake_in = du.input_data(sess)
    X_real_name = X_real[1]
    X_real = X_real[0]
    X_fake_name = X_fake_in[1]
    X_fake_in = X_fake_in[0]
    #X, Y = create_datasets(no_dir, yes_dir)
    #X_fake_in = tf.placeholder(tf.float32, (batch_size, 96, 96, 3))
    #X_real = tf.placeholder(tf.float32, (batch_size, 96, 96, 3))

    #######################
    # Instantiate generator
    #######################
    list_filters = [256, 128, 64, 3]
    list_strides = [2] * len(list_filters)
    list_kernel_size = [3] * len(list_filters)
    list_padding = ["SAME"] * len(list_filters)
    output_shape = X_real.get_shape().as_list()[1:]
    G = models.Generator(list_filters,
                         list_kernel_size,
                         list_strides,
                         list_padding,
                         output_shape,
                         batch_size=FLAGS.batch_size,
                         data_format=FLAGS.data_format)

    ###########################
    # Instantiate discriminator
    ###########################
    list_filters = [32, 64, 128, 256]
    list_strides = [2] * len(list_filters)
    list_kernel_size = [3] * len(list_filters)
    list_padding = ["SAME"] * len(list_filters)
    D = models.Discriminator(list_filters,
                             list_kernel_size,
                             list_strides,
                             list_padding,
                             FLAGS.batch_size,
                             data_format=FLAGS.data_format)

    ###########################
    # Instantiate optimizers
    ###########################
    G_opt = tf.train.AdamOptimizer(learning_rate=1E-4,
                                   name='G_opt',
                                   beta1=0.5,
                                   beta2=0.9)
    D_opt = tf.train.AdamOptimizer(learning_rate=1E-4,
                                   name='D_opt',
                                   beta1=0.5,
                                   beta2=0.9)

    ###########################
    # Instantiate model outputs
    ###########################

    # noise_input = tf.random_normal((FLAGS.batch_size, FLAGS.noise_dim,), stddev=0.1)
    noise_input = tf.random_uniform((
        FLAGS.batch_size,
        FLAGS.noise_dim,
    ),
                                    minval=-1,
                                    maxval=1)
    X_fake = G(X_fake_in)
    #X_fake = G(noise_input)

    # output images
    #X_G_input = du.unnormalize_image(X_fake_in)
    X_G_output = du.unnormalize_image(X_fake)
    X_real_output = du.unnormalize_image(X_real)

    D_real = D(X_real)
    D_fake = D(X_fake, reuse=True)

    ###########################
    # Instantiate losses
    ###########################

    G_loss = -tf.reduce_mean(D_fake) + (tf.reduce_mean(abs(X_fake - X_real)))
    D_loss = tf.reduce_mean(D_fake) - tf.reduce_mean(D_real)

    epsilon = tf.random_uniform(shape=[FLAGS.batch_size, 1, 1, 1],
                                minval=0.,
                                maxval=1.)
    X_hat = X_real + epsilon * (X_fake - X_real)
    D_X_hat = D(X_hat, reuse=True)
    grad_D_X_hat = tf.gradients(D_X_hat, [X_hat])[0]
    if FLAGS.data_format == "NCHW":
        red_idx = [1]
    else:
        red_idx = [-1]
    slopes = tf.sqrt(
        tf.reduce_sum(tf.square(grad_D_X_hat), reduction_indices=red_idx))
    gradient_penalty = tf.reduce_mean((slopes - 1.)**2)
    D_loss += 10 * gradient_penalty

    ###########################
    # Compute gradient updates
    ###########################

    dict_G_vars = G.get_trainable_variables()
    G_vars = [dict_G_vars[k] for k in dict_G_vars.keys()]

    dict_D_vars = D.get_trainable_variables()
    D_vars = [dict_D_vars[k] for k in dict_D_vars.keys()]

    G_gradvar = G_opt.compute_gradients(G_loss,
                                        var_list=G_vars,
                                        colocate_gradients_with_ops=True)
    G_update = G_opt.apply_gradients(G_gradvar, name='G_loss_minimize')

    D_gradvar = D_opt.compute_gradients(D_loss,
                                        var_list=D_vars,
                                        colocate_gradients_with_ops=True)
    D_update = D_opt.apply_gradients(D_gradvar, name='D_loss_minimize')

    ##########################
    # Group training ops
    ##########################
    loss_ops = [G_loss, D_loss]

    ##########################
    # Summary ops
    ##########################

    # Add summary for gradients
    tu.add_gradient_summary(G_gradvar)
    tu.add_gradient_summary(D_gradvar)

    # Add scalar symmaries
    tf.summary.scalar("G loss", G_loss)
    tf.summary.scalar("D loss", D_loss)
    tf.summary.scalar("gradient_penalty", gradient_penalty)

    summary_op = tf.summary.merge_all()

    ############################
    # Start training
    ############################

    # Initialize session
    saver = tu.initialize_session(sess)

    # Start queues
    tu.manage_queues(sess)

    # Summaries
    writer = tu.manage_summaries(sess)
    txtfile = open("/home/ubuntu/makeup_removal/WGAN-GP/testfile.txt", "w")
    g_loss_list = []
    d_loss_list = []
    for e in tqdm(range(FLAGS.nb_epoch), desc="Training progress"):

        t = tqdm(range(FLAGS.nb_batch_per_epoch),
                 desc="Epoch %i" % e,
                 mininterval=0.5)
        num = 0
        for batch_counter in t:
            #with_makeup_batch, without_makeup_batch = nextBatch(X, Y, num, batch_size)
            num += 1
            g_loss_total = 0
            d_loss_total = 0
            for di in range(5):
                sess.run([D_update])
            #    #output = sess.run([G_update] + loss_ops + [summary_op])

            #sess.run([D_update])
            output = sess.run([G_update] + loss_ops + [summary_op])
            g_loss, d_loss = sess.run([G_loss, D_loss])
            g_loss_total += g_loss
            d_loss_total += d_loss
            #print (g_loss, d_loss)

            if batch_counter % (FLAGS.nb_batch_per_epoch // 20) == 0:
                writer.add_summary(
                    output[-1], e * FLAGS.nb_batch_per_epoch + batch_counter)

            t.set_description('Epoch %i' % e)
        g_loss_list.append(g_loss_total)
        d_loss_list.append(d_loss_total)
        # Plot some generated images
        #output = sess.run([X_G_output, X_G_input, X_real_output])
        output = sess.run(
            [X_G_output, X_real_output, X_fake_name, X_real_name])
        vu.save_image(output[:2], FLAGS.data_format, e)
        name = output[2:]
        print(name)
        with open('/home/ubuntu/makeup_removal/WGAN-GP/testfile.txt',
                  'a') as f:
            f.write(str(e))
            f.write('\n')
            for array in name:
                for n in array:
                    f.write(str(n))
                    f.write('\n')
            f.write('\n\n')
        #txtfile.close()
        #print (name)

        # Save session
        saver.save(sess, os.path.join(FLAGS.model_dir, "model"), global_step=e)

    print('Finished training!')
    plt.plot(range(FLAGS.nb_epoch), g_loss_list, label='g_loss')
    plt.plot(range(FLAGS.nb_epoch), d_loss_list, label='d_loss')
    plt.legend()
    plt.title('G:1E-4, D:1E-4')
    plt.xlabel('epoch')
    plt.ylabel('loss')
    plt.savefig('/home/ubuntu/makeup_removal/WGAN-GP/G_4_D_4.png')
Esempio n. 6
0
def train_model():

    # Setup session
    sess = tu.setup_session()

    # Placeholder for data and Mnist iterator
    mnist = input_data.read_data_sets(FLAGS.raw_dir, one_hot=True)
    # images = tf.constant(mnist.train.images)
    if FLAGS.data_format == "NHWC":
        X_real = tf.placeholder(tf.float32,
                                shape=[FLAGS.batch_size, 28, 28, 1])
    else:
        X_real = tf.placeholder(tf.float32,
                                shape=[FLAGS.batch_size, 1, 28, 28])

    with tf.device('/cpu:0'):
        imgs = mnist.train.images.astype(np.float32)
        npts = imgs.shape[0]
        if FLAGS.data_format == "NHWC":
            imgs = imgs.reshape((npts, 28, 28, 1))
        else:
            imgs = imgs.reshape((npts, 1, 28, 28))
        imgs = (imgs - 0.5) / 0.5
        # input_images = tf.constant(imgs)

    #     image = tf.train.slice_input_producer([input_images], num_epochs=FLAGS.nb_epoch)
    #     X_real = tf.train.batch(image, batch_size=FLAGS.batch_size, num_threads=8)

    #######################
    # Instantiate generator
    #######################
    list_filters = [256, 1]
    list_strides = [2] * len(list_filters)
    list_kernel_size = [3] * len(list_filters)
    list_padding = ["SAME"] * len(list_filters)
    output_shape = X_real.get_shape().as_list()[1:]
    G = models.Generator(list_filters,
                         list_kernel_size,
                         list_strides,
                         list_padding,
                         output_shape,
                         batch_size=FLAGS.batch_size,
                         dset="mnist",
                         data_format=FLAGS.data_format)

    ###########################
    # Instantiate discriminator
    ###########################
    list_filters = [32, 64]
    list_strides = [2] * len(list_filters)
    list_kernel_size = [3] * len(list_filters)
    list_padding = ["SAME"] * len(list_filters)
    D = models.Discriminator(list_filters,
                             list_kernel_size,
                             list_strides,
                             list_padding,
                             FLAGS.batch_size,
                             data_format=FLAGS.data_format)

    ###########################
    # Instantiate optimizers
    ###########################
    G_opt = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate,
                                   name='G_opt',
                                   beta1=0.5)
    D_opt = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate,
                                   name='D_opt',
                                   beta1=0.5)

    ###########################
    # Instantiate model outputs
    ###########################

    # noise_input = tf.random_normal((FLAGS.batch_size, FLAGS.noise_dim,), stddev=0.1)
    noise_input = tf.random_uniform((
        FLAGS.batch_size,
        FLAGS.noise_dim,
    ),
                                    minval=-1,
                                    maxval=1)
    X_fake = G(noise_input)

    # output images
    X_G_output = du.unnormalize_image(X_fake)
    X_real_output = du.unnormalize_image(X_real)

    D_real = D(X_real)
    D_fake = D(X_fake, reuse=True)

    ###########################
    # Instantiate losses
    ###########################

    G_loss = objectives.binary_cross_entropy_with_logits(
        D_fake, tf.ones_like(D_fake))
    D_loss_real = objectives.binary_cross_entropy_with_logits(
        D_real, tf.ones_like(D_real))
    D_loss_fake = objectives.binary_cross_entropy_with_logits(
        D_fake, tf.zeros_like(D_fake))

    D_loss = D_loss_real + D_loss_fake

    # ######################################################################
    # # Some parameters need to be updated (e.g. BN moving average/variance)
    # ######################################################################
    # from tensorflow.python.ops import control_flow_ops
    # update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    # with tf.control_dependencies(update_ops):
    #     barrier = tf.no_op(name='update_barrier')
    # D_loss = control_flow_ops.with_dependencies([barrier], D_loss)
    # G_loss = control_flow_ops.with_dependencies([barrier], G_loss)

    ###########################
    # Compute gradient updates
    ###########################

    dict_G_vars = G.get_trainable_variables()
    G_vars = [dict_G_vars[k] for k in dict_G_vars.keys()]

    dict_D_vars = D.get_trainable_variables()
    D_vars = [dict_D_vars[k] for k in dict_D_vars.keys()]

    G_gradvar = G_opt.compute_gradients(G_loss, var_list=G_vars)
    G_update = G_opt.apply_gradients(G_gradvar, name='G_loss_minimize')

    D_gradvar = D_opt.compute_gradients(D_loss, var_list=D_vars)
    D_update = D_opt.apply_gradients(D_gradvar, name='D_loss_minimize')

    ##########################
    # Group training ops
    ##########################
    train_ops = [G_update, D_update]
    loss_ops = [G_loss, D_loss, D_loss_real, D_loss_fake]

    ##########################
    # Summary ops
    ##########################

    # Add summary for gradients
    tu.add_gradient_summary(G_gradvar)
    tu.add_gradient_summary(D_gradvar)

    # Add scalar symmaries
    tf.summary.scalar("G loss", G_loss)
    tf.summary.scalar("D loss real", D_loss_real)
    tf.summary.scalar("D loss fake", D_loss_fake)

    summary_op = tf.summary.merge_all()

    ############################
    # Start training
    ############################

    # Initialize session
    saver = tu.initialize_session(sess)

    # Start queues
    coord = tu.manage_queues(sess)

    # Summaries
    writer = tu.manage_summaries(sess)

    for e in tqdm(range(FLAGS.nb_epoch), desc="\nTraining progress"):

        t = tqdm(range(FLAGS.nb_batch_per_epoch),
                 desc="Epoch %i" % e,
                 mininterval=0.5)
        for batch_counter in t:

            # X_batch, _ = mnist.train.next_batch(FLAGS.batch_size)
            # if FLAGS.data_format == "NHWC":
            #     X_batch = np.reshape(X_batch, [-1, 28, 28, 1])
            # else:
            #     X_batch = np.reshape(X_batch, [-1, 1, 28, 28])
            # X_batch = (X_batch - 0.5) / 0.5

            X_batch = du.sample_batch(imgs, FLAGS.batch_size)
            output = sess.run(train_ops + loss_ops + [summary_op],
                              feed_dict={X_real: X_batch})

            if batch_counter % (FLAGS.nb_batch_per_epoch // 20) == 0:
                writer.add_summary(
                    output[-1], e * FLAGS.nb_batch_per_epoch + batch_counter)
            lossG, lossDreal, lossDfake = [output[2], output[4], output[5]]

            t.set_description(
                'Epoch %i: - G loss: %.2f D loss real: %.2f Dloss fake: %.2f' %
                (e, lossG, lossDreal, lossDfake))

            # variables = [v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)]
            # bmean = [v for v in variables if v.name == "generator/conv2D_1_1/BatchNorm/moving_mean:0"]
            # print sess.run(bmean)
            # raw_input()

        # Plot some generated images
        # output = sess.run([X_G_output, X_real_output])
        output = sess.run([X_G_output, X_real_output],
                          feed_dict={X_real: X_batch})
        vu.save_image(output, FLAGS.data_format, e)

        # Save session
        saver.save(sess, os.path.join(FLAGS.model_dir, "model"), global_step=e)

    print('Finished training!')
def train_model():

    # Setup session
    sess = tu.setup_session()

    # Setup async input queue of real images
    X_real = du.input_data(sess)

    #######################
    # Instantiate generator
    #######################
    list_filters = [256, 128, 64, 3]
    list_strides = [2] * len(list_filters)
    list_kernel_size = [3] * len(list_filters)
    list_padding = ["SAME"] * len(list_filters)
    output_shape = X_real.get_shape().as_list()[1:]
    G = models.Generator(list_filters,
                         list_kernel_size,
                         list_strides,
                         list_padding,
                         output_shape,
                         batch_size=FLAGS.batch_size,
                         data_format=FLAGS.data_format)

    ###########################
    # Instantiate discriminator
    ###########################
    list_filters = [32, 64, 128, 256]
    list_strides = [2] * len(list_filters)
    list_kernel_size = [3] * len(list_filters)
    list_padding = ["SAME"] * len(list_filters)
    D = models.Discriminator(list_filters,
                             list_kernel_size,
                             list_strides,
                             list_padding,
                             FLAGS.batch_size,
                             data_format=FLAGS.data_format)

    ###########################
    # Instantiate optimizers
    ###########################
    G_opt = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate,
                                   name='G_opt',
                                   beta1=0.5)
    D_opt = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate,
                                   name='D_opt',
                                   beta1=0.5)

    ###########################
    # Instantiate model outputs
    ###########################

    # noise_input = tf.random_normal((FLAGS.batch_size, FLAGS.noise_dim,), stddev=0.1, name="noise_input")
    noise_input = tf.random_uniform((
        FLAGS.batch_size,
        FLAGS.noise_dim,
    ),
                                    minval=-1,
                                    maxval=1,
                                    name="noise_input")
    X_fake = G(noise_input)

    # output images
    X_G_output = du.unnormalize_image(X_fake, name="X_G_output")
    X_real_output = du.unnormalize_image(X_real, name="X_real_output")

    D_real = D(X_real)
    D_fake = D(X_fake, reuse=True)

    ###########################
    # Instantiate losses
    ###########################

    G_loss = objectives.binary_cross_entropy_with_logits(
        D_fake, tf.ones_like(D_fake))
    D_loss_real = objectives.binary_cross_entropy_with_logits(
        D_real, tf.ones_like(D_real))
    D_loss_fake = objectives.binary_cross_entropy_with_logits(
        D_fake, tf.zeros_like(D_fake))

    # G_loss = objectives.wasserstein(D_fake, -tf.ones_like(D_fake))
    # D_loss_real = objectives.wasserstein(D_real, -tf.ones_like(D_real))
    # D_loss_fake = objectives.wasserstein(D_fake, tf.ones_like(D_fake))

    D_loss = D_loss_real + D_loss_fake

    ###########################
    # Compute gradient updates
    ###########################

    dict_G_vars = G.get_trainable_variables()
    G_vars = [dict_G_vars[k] for k in dict_G_vars.keys()]

    dict_D_vars = D.get_trainable_variables()
    D_vars = [dict_D_vars[k] for k in dict_D_vars.keys()]

    G_gradvar = G_opt.compute_gradients(G_loss, var_list=G_vars)
    G_update = G_opt.apply_gradients(G_gradvar, name='G_loss_minimize')

    D_gradvar = D_opt.compute_gradients(D_loss, var_list=D_vars)
    D_update = D_opt.apply_gradients(D_gradvar, name='D_loss_minimize')

    # clip_op = [var.assign(tf.clip_by_value(var, -0.01, 0.01)) for var in D_vars]

    G_update = G_opt.minimize(G_loss, var_list=G_vars, name='G_loss_minimize')
    D_update = D_opt.minimize(D_loss, var_list=D_vars, name='D_loss_minimize')

    ##########################
    # Group training ops
    ##########################
    train_ops = [G_update, D_update]
    loss_ops = [G_loss, D_loss, D_loss_real, D_loss_fake]

    ##########################
    # Summary ops
    ##########################

    # Add summary for gradients
    tu.add_gradient_summary(G_gradvar)
    tu.add_gradient_summary(D_gradvar)

    # Add scalar symmaries
    tf.summary.scalar("G loss", G_loss)
    tf.summary.scalar("D loss real", D_loss_real)
    tf.summary.scalar("D loss fake", D_loss_fake)

    summary_op = tf.summary.merge_all()

    ############################
    # Start training
    ############################

    # Initialize session
    saver = tu.initialize_session(sess)

    # Start queues
    coord = tu.manage_queues(sess)

    # Summaries
    writer = tu.manage_summaries(sess)

    # Run checks on data dimensions
    list_data = [noise_input, X_real, X_fake, X_G_output, X_real_output]
    output = sess.run([noise_input, X_real, X_fake, X_G_output, X_real_output])
    tu.check_data(output, list_data)

    for e in tqdm(range(FLAGS.nb_epoch), desc="Training progress"):

        list_G_loss = []
        list_D_loss_real = []
        list_D_loss_fake = []

        t = tqdm(range(FLAGS.nb_batch_per_epoch),
                 desc="Epoch %i" % e,
                 mininterval=0.5)
        for batch_counter in t:

            o_D = sess.run([D_update, D_loss_real, D_loss_fake])
            sess.run([G_update, G_loss])
            o_G = sess.run([G_update, G_loss])
            output = sess.run([summary_op])

            list_G_loss.append(o_G[-1])
            list_D_loss_real.append(o_D[-2])
            list_D_loss_fake.append(o_D[-1])

            # output = sess.run(train_ops + loss_ops + [summary_op])
            # list_G_loss.append(output[2])
            # list_D_loss_real.append(output[4])
            # list_D_loss_fake.append(output[5])

            if batch_counter % (FLAGS.nb_batch_per_epoch //
                                (int(0.5 * FLAGS.nb_batch_per_epoch))) == 0:
                writer.add_summary(
                    output[-1], e * FLAGS.nb_batch_per_epoch + batch_counter)

        t.set_description(
            'Epoch %i: - G loss: %.3f D loss real: %.3f Dloss fake: %.3f' %
            (e, np.mean(list_G_loss), np.mean(list_D_loss_real),
             np.mean(list_D_loss_fake)))

        # Plot some generated images
        output = sess.run([X_G_output, X_real_output])
        vu.save_image(output, FLAGS.data_format, e)

        # Save session
        saver.save(sess, os.path.join(FLAGS.model_dir, "model"), global_step=e)

        if e == 0:
            print(len(list_data))
            output = sess.run(
                [noise_input, X_real, X_fake, X_G_output, X_real_output])
            tu.check_data(output, list_data)

    print('Finished training!')
Esempio n. 8
0
def train_model():

    # Setup session
    sess = tu.setup_training_session()

    ##########
    # Innputs
    ##########

    # Setup async input queue of real images
    X_real = du.read_celebA()

    # Noise
    batch_size = tf.shape(X_real)[0]
    z_noise = tf.random_uniform((batch_size, FLAGS.z_dim),
                                minval=-1,
                                maxval=1,
                                name="z_input")
    epsilon = tf.random_uniform((batch_size, 1, 1, 1),
                                minval=0,
                                maxval=1,
                                name="epsilon")

    # learning rate
    lr_D = tf.Variable(initial_value=FLAGS.learning_rate,
                       trainable=False,
                       name='learning_rate')
    lr_G = tf.Variable(initial_value=FLAGS.learning_rate,
                       trainable=False,
                       name='learning_rate')

    ########################
    # Instantiate models
    ########################
    G = models.Generator()
    D = models.Discriminator()

    ###########################
    # Instantiate optimizers
    ###########################
    G_opt = tf.train.AdamOptimizer(learning_rate=lr_D, name='G_opt', beta1=0.5)
    D_opt = tf.train.AdamOptimizer(learning_rate=lr_G, name='D_opt', beta1=0.5)

    ##########
    # Outputs
    ##########
    X_fake = G(z_noise)
    X_hat = epsilon * X_real + (1 - epsilon) * X_fake

    D_real = D(X_real)
    D_fake = D(X_fake, reuse=True)
    D_X_hat = D(X_hat, reuse=True)

    grad_D_X_hat = tf.gradients(D_X_hat, X_hat)[0]

    # output images
    generated_toplot = du.unnormalize_image(X_fake, name="generated_toplot")
    real_toplot = du.unnormalize_image(X_real, name="real_toplot")

    ###########################
    # losses
    ###########################

    G_loss = losses.wasserstein(D_fake, -tf.ones_like(D_fake))
    D_loss_grad = FLAGS.lbd * tf.square((tf.nn.l2_loss(grad_D_X_hat) - 1))
    D_loss_real = losses.wasserstein(D_real, -tf.ones_like(D_real))
    D_loss_fake = losses.wasserstein(D_fake, tf.ones_like(D_fake))
    D_loss = D_loss_grad + D_loss_real + D_loss_fake

    ###########################
    # Compute updates ops
    ###########################

    dict_G_vars = G.get_trainable_variables()
    G_vars = [dict_G_vars[k] for k in dict_G_vars.keys()]

    dict_D_vars = D.get_trainable_variables()
    D_vars = [dict_D_vars[k] for k in dict_D_vars.keys()]

    G_gradvar = G_opt.compute_gradients(G_loss, var_list=G_vars)
    G_update = G_opt.apply_gradients(G_gradvar, name='G_loss_minimize')

    D_gradvar = D_opt.compute_gradients(D_loss, var_list=D_vars)
    D_update = D_opt.apply_gradients(D_gradvar, name='D_loss_minimize')

    # D_gradvar_fake = D_opt.compute_gradients(D_loss_fake, var_list=D_vars)
    # D_update_fake = D_opt.apply_gradients(D_gradvar_fake, name='D_loss_minimize_fake')

    ##########################
    # Summary ops
    ##########################

    # Add summary for gradients
    tu.add_gradient_summary(G_gradvar)
    tu.add_gradient_summary(D_gradvar)
    # tu.add_gradient_summary(D_gradvar_fake)

    # Add scalar symmaries for G
    tf.summary.scalar("G loss", G_loss)
    # Add scalar symmaries for D
    tf.summary.scalar("D loss real", D_loss_real)
    tf.summary.scalar("D loss fake", D_loss_fake)
    tf.summary.scalar("D loss grad", D_loss_grad)
    # Add scalar symmaries for D

    summary_op = tf.summary.merge_all()

    ############################
    # Start training
    ############################

    # Initialize session
    saver = tu.initialize_session(sess)

    # Start queues
    coord, threads = du.manage_queues(sess)

    # Summaries
    writer = tu.manage_summaries(sess)

    # Run checks on data dimensions
    list_data = [z_noise]
    list_data += [X_real, X_fake]
    list_data += [generated_toplot, real_toplot]
    output = sess.run(list_data)
    tu.check_data(output, list_data)

    for e in tqdm(range(FLAGS.nb_epoch), desc="Training progress"):

        t = tqdm(range(FLAGS.nb_batch_per_epoch),
                 desc="Epoch %i" % e,
                 mininterval=0.5)
        for batch_counter in t:

            # Update discriminator
            for i_D in range(FLAGS.ncritic):
                sess.run([D_update])
                # r = np.random.randint(0, 2)
                # if r == 0:
                #     sess.run([D_update_real])
                # else:
                #     sess.run([D_update_fake])

            # Update generator
            sess.run([G_update])

            if batch_counter % (FLAGS.nb_batch_per_epoch //
                                (int(0.5 * FLAGS.nb_batch_per_epoch))) == 0:
                output = sess.run([summary_op])
                writer.add_summary(
                    output[-1], e * FLAGS.nb_batch_per_epoch + batch_counter)

            t.set_description('Epoch %s:' % e)

        # Plot some generated images
        Xf, Xr = sess.run([generated_toplot, real_toplot])
        vu.save_image(Xf, Xr, title="current_batch", e=e)

        # Save session
        saver.save(sess, os.path.join(FLAGS.model_dir, "model"), global_step=e)

        # Show data statistics
        output = sess.run(list_data)
        tu.check_data(output, list_data)

    # Stop threads
    coord.request_stop()
    coord.join(threads)

    print('Finished training!')
Esempio n. 9
0
def train_model():

    # Setup session
    sess = tu.setup_session()

    # Setup async input queue of real images
    X_real = du.input_data(sess)

    #######################
    # Instantiate generator
    #######################
    list_filters = [256, 128, 64, 3]
    list_strides = [2] * len(list_filters)
    list_kernel_size = [3] * len(list_filters)
    list_padding = ["SAME"] * len(list_filters)
    output_shape = X_real.get_shape().as_list()[1:]
    G = models.Generator(list_filters,
                         list_kernel_size,
                         list_strides,
                         list_padding,
                         output_shape,
                         batch_size=FLAGS.batch_size,
                         data_format=FLAGS.data_format)

    ###########################
    # Instantiate discriminator
    ###########################
    list_filters = [32, 64, 128, 256]
    list_strides = [2] * len(list_filters)
    list_kernel_size = [3] * len(list_filters)
    list_padding = ["SAME"] * len(list_filters)
    D = models.Discriminator(list_filters,
                             list_kernel_size,
                             list_strides,
                             list_padding,
                             FLAGS.batch_size,
                             data_format=FLAGS.data_format)

    ###########################
    # Instantiate optimizers
    ###########################
    G_opt = tf.train.AdamOptimizer(learning_rate=1E-4,
                                   name='G_opt',
                                   beta1=0.5,
                                   beta2=0.9)
    D_opt = tf.train.AdamOptimizer(learning_rate=1E-4,
                                   name='D_opt',
                                   beta1=0.5,
                                   beta2=0.9)

    ###########################
    # Instantiate model outputs
    ###########################

    # noise_input = tf.random_normal((FLAGS.batch_size, FLAGS.noise_dim,), stddev=0.1)
    noise_input = tf.random_uniform((
        FLAGS.batch_size,
        FLAGS.noise_dim,
    ),
                                    minval=-1,
                                    maxval=1)
    X_fake = G(noise_input)

    # output images
    X_G_output = du.unnormalize_image(X_fake)
    X_real_output = du.unnormalize_image(X_real)

    D_real = D(X_real)
    D_fake = D(X_fake, reuse=True)

    ###########################
    # Instantiate losses
    ###########################

    G_loss = -tf.reduce_mean(D_fake)
    D_loss = tf.reduce_mean(D_fake) - tf.reduce_mean(D_real)

    epsilon = tf.random_uniform(shape=[FLAGS.batch_size, 1, 1, 1],
                                minval=0.,
                                maxval=1.)
    X_hat = X_real + epsilon * (X_fake - X_real)
    D_X_hat = D(X_hat, reuse=True)
    grad_D_X_hat = tf.gradients(D_X_hat, [X_hat])[0]
    if FLAGS.data_format == "NCHW":
        red_idx = [1]
    else:
        red_idx = [-1]
    slopes = tf.sqrt(
        tf.reduce_sum(tf.square(grad_D_X_hat), reduction_indices=red_idx))
    gradient_penalty = tf.reduce_mean((slopes - 1.)**2)
    D_loss += 10 * gradient_penalty

    ###########################
    # Compute gradient updates
    ###########################

    dict_G_vars = G.get_trainable_variables()
    G_vars = [dict_G_vars[k] for k in dict_G_vars.keys()]

    dict_D_vars = D.get_trainable_variables()
    D_vars = [dict_D_vars[k] for k in dict_D_vars.keys()]

    G_gradvar = G_opt.compute_gradients(G_loss,
                                        var_list=G_vars,
                                        colocate_gradients_with_ops=True)
    G_update = G_opt.apply_gradients(G_gradvar, name='G_loss_minimize')

    D_gradvar = D_opt.compute_gradients(D_loss,
                                        var_list=D_vars,
                                        colocate_gradients_with_ops=True)
    D_update = D_opt.apply_gradients(D_gradvar, name='D_loss_minimize')

    ##########################
    # Group training ops
    ##########################
    loss_ops = [G_loss, D_loss]

    ##########################
    # Summary ops
    ##########################

    # Add summary for gradients
    tu.add_gradient_summary(G_gradvar)
    tu.add_gradient_summary(D_gradvar)

    # Add scalar symmaries
    tf.summary.scalar("G loss", G_loss)
    tf.summary.scalar("D loss", D_loss)
    tf.summary.scalar("gradient_penalty", gradient_penalty)

    summary_op = tf.summary.merge_all()

    ############################
    # Start training
    ############################

    # Initialize session
    saver = tu.initialize_session(sess)

    # Start queues
    tu.manage_queues(sess)

    # Summaries
    writer = tu.manage_summaries(sess)

    for e in tqdm(range(FLAGS.nb_epoch), desc="Training progress"):

        t = tqdm(range(FLAGS.nb_batch_per_epoch),
                 desc="Epoch %i" % e,
                 mininterval=0.5)
        for batch_counter in t:

            for di in range(5):
                sess.run([D_update])

            output = sess.run([G_update] + loss_ops + [summary_op])

            if batch_counter % (FLAGS.nb_batch_per_epoch // 20) == 0:
                writer.add_summary(
                    output[-1], e * FLAGS.nb_batch_per_epoch + batch_counter)

            t.set_description('Epoch %i' % e)

        # Plot some generated images
        output = sess.run([X_G_output, X_real_output])
        vu.save_image(output, FLAGS.data_format, e)

        # Save session
        saver.save(sess, os.path.join(FLAGS.model_dir, "model"), global_step=e)

    print('Finished training!')
def train_model():

    # Setup session
    sess = tu.setup_session()

    # Setup async input queue of real images
    X_real = du.input_data(sess)

    #######################
    # Instantiate generator
    #######################
    list_filters = [256, 128, 64, 3]
    list_strides = [2] * len(list_filters)
    list_kernel_size = [3] * len(list_filters)
    list_padding = ["SAME"] * len(list_filters)
    output_shape = X_real.get_shape().as_list()[1:]
    G = models.Generator(list_filters, list_kernel_size, list_strides, list_padding, output_shape,
                         batch_size=FLAGS.batch_size, data_format=FLAGS.data_format)

    ###########################
    # Instantiate discriminator
    ###########################
    list_filters = [32, 64, 128, 256]
    list_strides = [2] * len(list_filters)
    list_kernel_size = [3] * len(list_filters)
    list_padding = ["SAME"] * len(list_filters)
    D = models.Discriminator(list_filters, list_kernel_size, list_strides, list_padding,
                             FLAGS.batch_size, data_format=FLAGS.data_format)

    ###########################
    # Instantiate optimizers
    ###########################
    G_opt = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate, name='G_opt', beta1=0.5)
    D_opt = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate, name='D_opt', beta1=0.5)

    ###########################
    # Instantiate model outputs
    ###########################

    # noise_input = tf.random_normal((FLAGS.batch_size, FLAGS.noise_dim,), stddev=0.1, name="noise_input")
    noise_input = tf.random_uniform((FLAGS.batch_size, FLAGS.noise_dim,), minval=-1, maxval=1, name="noise_input")
    X_fake = G(noise_input)

    # output images
    X_G_output = du.unnormalize_image(X_fake, name="X_G_output")
    X_real_output = du.unnormalize_image(X_real, name="X_real_output")

    D_real = D(X_real)
    D_fake = D(X_fake, reuse=True)

    ###########################
    # Instantiate losses
    ###########################

    G_loss = objectives.binary_cross_entropy_with_logits(D_fake, tf.ones_like(D_fake))
    D_loss_real = objectives.binary_cross_entropy_with_logits(D_real, tf.ones_like(D_real))
    D_loss_fake = objectives.binary_cross_entropy_with_logits(D_fake, tf.zeros_like(D_fake))

    # G_loss = objectives.wasserstein(D_fake, -tf.ones_like(D_fake))
    # D_loss_real = objectives.wasserstein(D_real, -tf.ones_like(D_real))
    # D_loss_fake = objectives.wasserstein(D_fake, tf.ones_like(D_fake))

    D_loss = D_loss_real + D_loss_fake

    ###########################
    # Compute gradient updates
    ###########################

    dict_G_vars = G.get_trainable_variables()
    G_vars = [dict_G_vars[k] for k in dict_G_vars.keys()]

    dict_D_vars = D.get_trainable_variables()
    D_vars = [dict_D_vars[k] for k in dict_D_vars.keys()]

    G_gradvar = G_opt.compute_gradients(G_loss, var_list=G_vars)
    G_update = G_opt.apply_gradients(G_gradvar, name='G_loss_minimize')

    D_gradvar = D_opt.compute_gradients(D_loss, var_list=D_vars)
    D_update = D_opt.apply_gradients(D_gradvar, name='D_loss_minimize')

    # clip_op = [var.assign(tf.clip_by_value(var, -0.01, 0.01)) for var in D_vars]

    G_update = G_opt.minimize(G_loss, var_list=G_vars, name='G_loss_minimize')
    D_update = D_opt.minimize(D_loss, var_list=D_vars, name='D_loss_minimize')

    ##########################
    # Group training ops
    ##########################
    train_ops = [G_update, D_update]
    loss_ops = [G_loss, D_loss, D_loss_real, D_loss_fake]

    ##########################
    # Summary ops
    ##########################

    # Add summary for gradients
    tu.add_gradient_summary(G_gradvar)
    tu.add_gradient_summary(D_gradvar)

    # Add scalar symmaries
    tf.summary.scalar("G loss", G_loss)
    tf.summary.scalar("D loss real", D_loss_real)
    tf.summary.scalar("D loss fake", D_loss_fake)

    summary_op = tf.summary.merge_all()

    ############################
    # Start training
    ############################

    # Initialize session
    saver = tu.initialize_session(sess)

    # Start queues
    coord = tu.manage_queues(sess)

    # Summaries
    writer = tu.manage_summaries(sess)

    # Run checks on data dimensions
    list_data = [noise_input, X_real, X_fake, X_G_output, X_real_output]
    output = sess.run([noise_input, X_real, X_fake, X_G_output, X_real_output])
    tu.check_data(output, list_data)

    for e in tqdm(range(FLAGS.nb_epoch), desc="Training progress"):

        list_G_loss = []
        list_D_loss_real = []
        list_D_loss_fake = []

        t = tqdm(range(FLAGS.nb_batch_per_epoch), desc="Epoch %i" % e, mininterval=0.5)
        for batch_counter in t:

            o_D = sess.run([D_update, D_loss_real, D_loss_fake])
            sess.run([G_update, G_loss])
            o_G = sess.run([G_update, G_loss])
            output = sess.run([summary_op])

            list_G_loss.append(o_G[-1])
            list_D_loss_real.append(o_D[-2])
            list_D_loss_fake.append(o_D[-1])

            # output = sess.run(train_ops + loss_ops + [summary_op])
            # list_G_loss.append(output[2])
            # list_D_loss_real.append(output[4])
            # list_D_loss_fake.append(output[5])

            if batch_counter % (FLAGS.nb_batch_per_epoch // (int(0.5 * FLAGS.nb_batch_per_epoch))) == 0:
                writer.add_summary(output[-1], e * FLAGS.nb_batch_per_epoch + batch_counter)

            t.set_description('Epoch %i: - G loss: %.3f D loss real: %.3f Dloss fake: %.3f' %
                              (e, np.mean(list_G_loss), np.mean(list_D_loss_real), np.mean(list_D_loss_fake)))

        # Plot some generated images
        output = sess.run([X_G_output, X_real_output])
        vu.save_image(output, FLAGS.data_format, e)

        # Save session
        saver.save(sess, os.path.join(FLAGS.model_dir, "model"), global_step=e)

        if e == 0:
            print len(list_data)
            output = sess.run([noise_input, X_real, X_fake, X_G_output, X_real_output])
            tu.check_data(output, list_data)

    print('Finished training!')
def train_model():

    # Setup session
    sess = tu.setup_session()

    # Setup async input queue of real images
    X_real = du.input_data(sess)

    #######################
    # Instantiate generator
    #######################
    list_filters = [256, 128, 64, 3]
    list_strides = [2] * len(list_filters)
    list_kernel_size = [3] * len(list_filters)
    list_padding = ["SAME"] * len(list_filters)
    output_shape = X_real.get_shape().as_list()[1:]
    G = models.Generator(list_filters, list_kernel_size, list_strides, list_padding, output_shape,
                         batch_size=FLAGS.batch_size, data_format=FLAGS.data_format)

    ###########################
    # Instantiate discriminator
    ###########################
    list_filters = [32, 64, 128, 256]
    list_strides = [2] * len(list_filters)
    list_kernel_size = [3] * len(list_filters)
    list_padding = ["SAME"] * len(list_filters)
    D = models.Discriminator(list_filters, list_kernel_size, list_strides, list_padding,
                             FLAGS.batch_size, data_format=FLAGS.data_format)

    ###########################
    # Instantiate optimizers
    ###########################
    G_opt = tf.train.AdamOptimizer(learning_rate=1E-4, name='G_opt', beta1=0.5, beta2=0.9)
    D_opt = tf.train.AdamOptimizer(learning_rate=1E-4, name='D_opt', beta1=0.5, beta2=0.9)

    ###########################
    # Instantiate model outputs
    ###########################

    # noise_input = tf.random_normal((FLAGS.batch_size, FLAGS.noise_dim,), stddev=0.1)
    noise_input = tf.random_uniform((FLAGS.batch_size, FLAGS.noise_dim,), minval=-1, maxval=1)
    X_fake = G(noise_input)

    # output images
    X_G_output = du.unnormalize_image(X_fake)
    X_real_output = du.unnormalize_image(X_real)

    D_real = D(X_real)
    D_fake = D(X_fake, reuse=True)

    ###########################
    # Instantiate losses
    ###########################

    G_loss = -tf.reduce_mean(D_fake)
    D_loss = tf.reduce_mean(D_fake) - tf.reduce_mean(D_real)

    epsilon = tf.random_uniform(
        shape=[FLAGS.batch_size, 1, 1, 1],
        minval=0.,
        maxval=1.
    )
    X_hat = X_real + epsilon * (X_fake - X_real)
    D_X_hat = D(X_hat, reuse=True)
    grad_D_X_hat = tf.gradients(D_X_hat, [X_hat])[0]
    if FLAGS.data_format == "NCHW":
        red_idx = [1]
    else:
        red_idx = [-1]
    slopes = tf.sqrt(tf.reduce_sum(tf.square(grad_D_X_hat), reduction_indices=red_idx))
    gradient_penalty = tf.reduce_mean((slopes - 1.)**2)
    D_loss += 10 * gradient_penalty

    ###########################
    # Compute gradient updates
    ###########################

    dict_G_vars = G.get_trainable_variables()
    G_vars = [dict_G_vars[k] for k in dict_G_vars.keys()]

    dict_D_vars = D.get_trainable_variables()
    D_vars = [dict_D_vars[k] for k in dict_D_vars.keys()]

    G_gradvar = G_opt.compute_gradients(G_loss, var_list=G_vars, colocate_gradients_with_ops=True)
    G_update = G_opt.apply_gradients(G_gradvar, name='G_loss_minimize')

    D_gradvar = D_opt.compute_gradients(D_loss, var_list=D_vars, colocate_gradients_with_ops=True)
    D_update = D_opt.apply_gradients(D_gradvar, name='D_loss_minimize')

    ##########################
    # Group training ops
    ##########################
    loss_ops = [G_loss, D_loss]

    ##########################
    # Summary ops
    ##########################

    # Add summary for gradients
    tu.add_gradient_summary(G_gradvar)
    tu.add_gradient_summary(D_gradvar)

    # Add scalar symmaries
    tf.summary.scalar("G loss", G_loss)
    tf.summary.scalar("D loss", D_loss)
    tf.summary.scalar("gradient_penalty", gradient_penalty)

    summary_op = tf.summary.merge_all()

    ############################
    # Start training
    ############################

    # Initialize session
    saver = tu.initialize_session(sess)

    # Start queues
    tu.manage_queues(sess)

    # Summaries
    writer = tu.manage_summaries(sess)

    for e in tqdm(range(FLAGS.nb_epoch), desc="Training progress"):

        t = tqdm(range(FLAGS.nb_batch_per_epoch), desc="Epoch %i" % e, mininterval=0.5)
        for batch_counter in t:

            for di in range(5):
                sess.run([D_update])

            output = sess.run([G_update] + loss_ops + [summary_op])

            if batch_counter % (FLAGS.nb_batch_per_epoch // 20) == 0:
                writer.add_summary(output[-1], e * FLAGS.nb_batch_per_epoch + batch_counter)

            t.set_description('Epoch %i' % e)

        # Plot some generated images
        output = sess.run([X_G_output, X_real_output])
        vu.save_image(output, FLAGS.data_format, e)

        # Save session
        saver.save(sess, os.path.join(FLAGS.model_dir, "model"), global_step=e)

    print('Finished training!')