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