def construct_latent()

in ebm_sandbox.py [0:0]


def construct_latent(weights, X, Y_GT, model, target_vars):

    eps = 0.001
    X_init = X[0:1]

    def traversals(model, X, weights, Y_GT):
        if FLAGS.hessian:
            e_pos = model.forward(X, weights, label=Y_GT)
            hessian = tf.hessians(e_pos, X)
            hessian = tf.reshape(hessian, (1, 64*64, 64*64))[0]
            e, v = tf.linalg.eigh(hessian)
        else:
            latent = model.forward(X, weights, label=Y_GT, return_logit=True)
            latents = tf.split(latent, 128, axis=1)
            jacobian = [tf.gradients(latent, X)[0] for latent in latents]
            jacobian = tf.stack(jacobian, axis=1)
            jacobian = tf.reshape(jacobian, (tf.shape(jacobian)[1], tf.shape(jacobian)[1], 64*64))
            s, _, v = tf.linalg.svd(jacobian)

        return v


    var_scale = 1.0
    n = 3
    xs = []

    v = traversals(model, X_init, weights, Y_GT)

    for i in range(n):
        var = tf.reshape(v[:, i], (1, 64, 64))
        X_plus = X_init - var_scale * var
        X_min = X_init + var_scale * var

        xs.extend([X_plus, X_min])

    x_stack = tf.stack(xs, axis=0)

    e_pos_hess_modify = model.forward(x_stack, weights, label=Y_GT)

    for i in range(20):
        x_stack = x_stack + tf.random_normal(tf.shape(x_stack), mean=0.0, stddev=0.005)
        e_pos = model.forward(x_stack, weights, label=Y_GT)

        x_grad = tf.gradients(e_pos, [x_stack])[0]
        x_stack = x_stack - 4*FLAGS.step_lr * x_grad

        x_stack = tf.clip_by_value(x_stack, 0, 1)

    x_mods = tf.split(X, 6)

    eigs = []
    for j in range(6):
        x_mod = x_mods[j]
        v = traversals(model, x_mod, weights, Y_GT)

        idx = j // 2
        var = tf.reshape(v[:, idx], (1, 64, 64))

        if j % 2 == 1:
            x_mod = x_mod + var_scale * var
            eigs.append(var)
        else:
            x_mod = x_mod - var_scale * var
            eigs.append(-var)

        x_mod = tf.clip_by_value(x_mod, 0, 1)
        x_mods[j] = x_mod

    x_mods_stack = tf.stack(x_mods, axis=0)

    eigs_stack = tf.stack(eigs, axis=0)
    energys = []

    for i in range(20):
        x_mods_stack = x_mods_stack + tf.random_normal(tf.shape(x_mods_stack), mean=0.0, stddev=0.005)
        e_pos = model.forward(x_mods_stack, weights, label=Y_GT)

        x_grad = tf.gradients(e_pos, [x_mods_stack])[0]
        x_mods_stack = x_mods_stack - 4*FLAGS.step_lr * x_grad
        # x_mods_stack = x_mods_stack + 0.1 * eigs_stack

        x_mods_stack = tf.clip_by_value(x_mods_stack, 0, 1)

        energys.append(e_pos)

    x_refine = x_mods_stack
    es = tf.stack(energys, axis=0)

    # target_vars['hessian'] = hessian
    # target_vars['e'] = e
    target_vars['v'] = v
    target_vars['x_stack'] = x_stack
    target_vars['x_refine'] = x_refine
    target_vars['es'] = es