def build_training_graph(x, y, ul_x, ul_u, lr, mom):
    global_step = tf.get_variable(name='global_step', shape=[], dtype=tf.float32, initializer=tf.constant_initializer(0.0), trainable=False)
    logit = vat.forward(x)
    nll_loss = L.ce_loss(logit, y)
    with tf.variable_scope(tf.get_variable_scope(), reuse=True):
        if (FLAGS.method == 'vat'):
            ul_logit = vat.forward(ul_x, is_training=True, update_batch_stats=False)
            (vat_loss, ul_u_updated) = vat.virtual_adversarial_loss(ul_x, ul_u, ul_logit)
            additional_loss = vat_loss
        elif (FLAGS.method == 'vatent'):
            ul_logit = vat.forward(ul_x, is_training=True, update_batch_stats=False)
            (vat_loss, ul_u_updated) = vat.virtual_adversarial_loss(ul_x, ul_u, ul_logit)
            ent_loss = L.entropy_y_x(ul_logit)
            additional_loss = (vat_loss + ent_loss)
        elif (FLAGS.method == 'baseline'):
            additional_loss = 0
        else:
            raise NotImplementedError
        loss = (nll_loss + additional_loss)
    opt = tf.train.AdamOptimizer(learning_rate=lr, beta1=mom)
    tvars = tf.trainable_variables()
    grads_and_vars = opt.compute_gradients(loss, tvars)
    train_op = opt.apply_gradients(grads_and_vars, global_step=global_step)
    return (loss, train_op, global_step, ul_u_updated)
def build_training_graph(x, y, ul_x, ul_u, lr, mom):
    global_step = tf.get_variable(name='global_step', shape=[], dtype=tf.float32, initializer=tf.constant_initializer(0.0), trainable=False)
    logit = vat.forward(x)
    nll_loss = L.ce_loss(logit, y)
    with tf.variable_scope(tf.get_variable_scope(), reuse=True):
        if (FLAGS.method == 'vat'):
            ul_logit = vat.forward(ul_x, is_training=True, update_batch_stats=False)
            (vat_loss, ul_u_updated) = vat.virtual_adversarial_loss(ul_x, ul_u, ul_logit)
            additional_loss = vat_loss
        elif (FLAGS.method == 'vatent'):
            ul_logit = vat.forward(ul_x, is_training=True, update_batch_stats=False)
            (vat_loss, ul_u_updated) = vat.virtual_adversarial_loss(ul_x, ul_u, ul_logit)
            ent_loss = L.entropy_y_x(ul_logit)
            additional_loss = (vat_loss + ent_loss)
        elif (FLAGS.method == 'baseline'):
            additional_loss = 0
        else:
            raise NotImplementedError
        loss = (nll_loss + additional_loss)
    opt = tf.train.AdamOptimizer(learning_rate=lr, beta1=mom)
    tvars = tf.trainable_variables()
    grads_and_vars = opt.compute_gradients(loss, tvars)
    train_op = opt.apply_gradients(grads_and_vars, global_step=global_step)
    return (loss, train_op, global_step, ul_u_updated)