in ebm_combine.py [0:0]
def gentest(sess, kvs, data, latents, save_exp_dir):
X_NOISE = kvs['X_NOISE']
LABEL_SIZE = kvs['LABEL_SIZE']
LABEL_SHAPE = kvs['LABEL_SHAPE']
LABEL_POS = kvs['LABEL_POS']
LABEL_ROT = kvs['LABEL_ROT']
model_size = kvs['model_size']
model_shape = kvs['model_shape']
model_pos = kvs['model_pos']
model_rot = kvs['model_rot']
weight_size = kvs['weight_size']
weight_shape = kvs['weight_shape']
weight_pos = kvs['weight_pos']
weight_rot = kvs['weight_rot']
X = tf.placeholder(shape=(None, 64, 64), dtype=tf.float32)
datafull = data
# Test combination of generalization where we use slices of both training
x_final = X_NOISE
x_mod_size = X_NOISE
x_mod_pos = X_NOISE
for i in range(FLAGS.num_steps):
# use cond_pos
energies = []
x_mod_pos = x_mod_pos + tf.random_normal(tf.shape(x_mod_pos), mean=0.0, stddev=0.005)
e_noise = model_pos.forward(x_final, weight_pos, label=LABEL_POS)
# energies.append(e_noise)
x_grad = tf.gradients(e_noise, [x_final])[0]
x_mod_pos = x_mod_pos + tf.random_normal(tf.shape(x_mod_pos), mean=0.0, stddev=0.005)
x_mod_pos = x_mod_pos - FLAGS.step_lr * x_grad
x_mod_pos = tf.clip_by_value(x_mod_pos, 0, 1)
if FLAGS.joint_shape:
# use cond_shape
e_noise = model_shape.forward(x_mod_pos, weight_shape, label=LABEL_SHAPE)
elif FLAGS.joint_rot:
e_noise = model_rot.forward(x_mod_pos, weight_rot, label=LABEL_ROT)
else:
# use cond_size
e_noise = model_size.forward(x_mod_pos, weight_size, label=LABEL_SIZE)
# energies.append(e_noise)
# energy_stack = tf.concat(energies, axis=1)
# energy_stack = tf.reduce_logsumexp(-1*energy_stack, axis=1)
# energy_stack = tf.reduce_sum(energy_stack, axis=1)
x_grad = tf.gradients(e_noise, [x_mod_pos])[0]
x_mod_pos = x_mod_pos - FLAGS.step_lr * x_grad
x_mod_pos = tf.clip_by_value(x_mod_pos, 0, 1)
# for x_mod_size
# use cond_size
# e_noise = model_size.forward(x_mod_size, weight_size, label=LABEL_SIZE)
# x_grad = tf.gradients(e_noise, [x_mod_size])[0]
# x_mod_size = x_mod_size + tf.random_normal(tf.shape(x_mod_size), mean=0.0, stddev=0.005)
# x_mod_size = x_mod_size - FLAGS.step_lr * x_grad
# x_mod_size = tf.clip_by_value(x_mod_size, 0, 1)
# # use cond_pos
# e_noise = model_pos.forward(x_mod_size, weight_pos, label=LABEL_POS)
# x_grad = tf.gradients(e_noise, [x_mod_size])[0]
# x_mod_size = x_mod_size + tf.random_normal(tf.shape(x_mod_size), mean=0.0, stddev=0.005)
# x_mod_size = x_mod_size - FLAGS.step_lr * tf.stop_gradient(x_grad)
# x_mod_size = tf.clip_by_value(x_mod_size, 0, 1)
x_mod = x_mod_pos
x_final = x_mod
if FLAGS.joint_shape:
loss_kl = model_shape.forward(x_final, weight_shape, reuse=True, label=LABEL_SHAPE, stop_grad=True) + \
model_pos.forward(x_final, weight_pos, reuse=True, label=LABEL_POS, stop_grad=True)
energy_pos = model_shape.forward(X, weight_shape, reuse=True, label=LABEL_SHAPE) + \
model_pos.forward(X, weight_pos, reuse=True, label=LABEL_POS)
energy_neg = model_shape.forward(tf.stop_gradient(x_mod), weight_shape, reuse=True, label=LABEL_SHAPE) + \
model_pos.forward(tf.stop_gradient(x_mod), weight_pos, reuse=True, label=LABEL_POS)
elif FLAGS.joint_rot:
loss_kl = model_rot.forward(x_final, weight_rot, reuse=True, label=LABEL_ROT, stop_grad=True) + \
model_pos.forward(x_final, weight_pos, reuse=True, label=LABEL_POS, stop_grad=True)
energy_pos = model_rot.forward(X, weight_rot, reuse=True, label=LABEL_ROT) + \
model_pos.forward(X, weight_pos, reuse=True, label=LABEL_POS)
energy_neg = model_rot.forward(tf.stop_gradient(x_mod), weight_rot, reuse=True, label=LABEL_ROT) + \
model_pos.forward(tf.stop_gradient(x_mod), weight_pos, reuse=True, label=LABEL_POS)
else:
loss_kl = model_size.forward(x_final, weight_size, reuse=True, label=LABEL_SIZE, stop_grad=True) + \
model_pos.forward(x_final, weight_pos, reuse=True, label=LABEL_POS, stop_grad=True)
energy_pos = model_size.forward(X, weight_size, reuse=True, label=LABEL_SIZE) + \
model_pos.forward(X, weight_pos, reuse=True, label=LABEL_POS)
energy_neg = model_size.forward(tf.stop_gradient(x_mod), weight_size, reuse=True, label=LABEL_SIZE) + \
model_pos.forward(tf.stop_gradient(x_mod), weight_pos, reuse=True, label=LABEL_POS)
energy_neg_reduced = (energy_neg - tf.reduce_min(energy_neg))
coeff = tf.stop_gradient(tf.exp(-energy_neg_reduced))
norm_constant = tf.stop_gradient(tf.reduce_sum(coeff)) + 1e-4
neg_loss = coeff * (-1*energy_neg) / norm_constant
loss_ml = tf.reduce_mean(energy_pos) - tf.reduce_mean(energy_neg)
loss_total = loss_ml + tf.reduce_mean(loss_kl) + 1 * (tf.reduce_mean(tf.square(energy_pos)) + tf.reduce_mean(tf.square(energy_neg)))
optimizer = AdamOptimizer(1e-3, beta1=0.0, beta2=0.999)
gvs = optimizer.compute_gradients(loss_total)
gvs = [(k, v) for (k, v) in gvs if k is not None]
train_op = optimizer.apply_gradients(gvs)
vs = optimizer.variables()
sess.run(tf.variables_initializer(vs))
dataloader = DataLoader(DSpritesGen(data, latents), batch_size=FLAGS.batch_size, num_workers=6, drop_last=True, shuffle=True)
x_off = tf.reduce_mean(tf.square(x_mod - X))
itr = 0
saver = tf.train.Saver()
x_mod = None
if FLAGS.train:
replay_buffer = ReplayBuffer(10000)
for _ in range(1):
for data_corrupt, data, label_size, label_pos in tqdm(dataloader):
data_corrupt = data_corrupt.numpy()[:, :, :]
data = data.numpy()[:, :, :]
if x_mod is not None:
replay_buffer.add(x_mod)
replay_batch = replay_buffer.sample(FLAGS.batch_size)
replay_mask = (np.random.uniform(0, 1, (FLAGS.batch_size)) > 0.95)
data_corrupt[replay_mask] = replay_batch[replay_mask]
if FLAGS.joint_shape:
feed_dict = {X_NOISE: data_corrupt, X: data, LABEL_SHAPE: label_size, LABEL_POS: label_pos}
elif FLAGS.joint_rot:
feed_dict = {X_NOISE: data_corrupt, X: data, LABEL_ROT: label_size, LABEL_POS: label_pos}
else:
feed_dict = {X_NOISE: data_corrupt, X: data, LABEL_SIZE: label_size, LABEL_POS: label_pos}
_, off_value, e_pos, e_neg, x_mod = sess.run([train_op, x_off, energy_pos, energy_neg, x_final], feed_dict=feed_dict)
itr += 1
if itr % 10 == 0:
print("x_off of {}, e_pos of {}, e_neg of {} itr of {}".format(off_value, e_pos.mean(), e_neg.mean(), itr))
if itr == FLAGS.break_steps:
break
saver.save(sess, osp.join(save_exp_dir, 'model_gentest'))
saver.restore(sess, osp.join(save_exp_dir, 'model_gentest'))
l = latents
if FLAGS.joint_shape:
mask_gen = (l[:, 3] == 30 * np.pi / 39) * (l[:, 2] == 0.5)
elif FLAGS.joint_rot:
mask_gen = (l[:, 1] == 1) * (l[:, 2] == 0.5)
else:
mask_gen = (l[:, 3] == 30 * np.pi / 39) * (l[:, 1] == 1) & (~((l[:, 2] == 0.5) | ((l[:, 4] == 16/31) & (l[:, 5] == 16/31))))
data_gen = datafull[mask_gen]
latents_gen = latents[mask_gen]
losses = []
for dat, latent in zip(np.array_split(data_gen, 120), np.array_split(latents_gen, 120)):
x = 0.5 + np.random.randn(*dat.shape)
if FLAGS.joint_shape:
feed_dict = {LABEL_SHAPE: np.eye(3)[latent[:, 1].astype(np.int32) - 1], LABEL_POS: latent[:, 4:], X_NOISE: x, X: dat}
elif FLAGS.joint_rot:
feed_dict = {LABEL_ROT: np.concatenate([np.cos(latent[:, 3:4]), np.sin(latent[:, 3:4])], axis=1), LABEL_POS: latent[:, 4:], X_NOISE: x, X: dat}
else:
feed_dict = {LABEL_SIZE: latent[:, 2:3], LABEL_POS: latent[:, 4:], X_NOISE: x, X: dat}
for i in range(2):
x = sess.run([x_final], feed_dict=feed_dict)[0]
feed_dict[X_NOISE] = x
loss = sess.run([x_off], feed_dict=feed_dict)[0]
losses.append(loss)
print("Mean MSE loss of {} ".format(np.mean(losses)))
data_try = data_gen[:10]
data_init = 0.5 + 0.5 * np.random.randn(10, 64, 64)
latent_scale = latents_gen[:10, 2:3]
latent_pos = latents_gen[:10, 4:]
if FLAGS.joint_shape:
feed_dict = {X_NOISE: data_init, LABEL_SHAPE: np.eye(3)[latent[:10, 1].astype(np.int32)-1], LABEL_POS: latent_pos}
elif FLAGS.joint_rot:
feed_dict = {LABEL_ROT: np.concatenate([np.cos(latent[:10, 3:4]), np.sin(latent[:10, 3:4])], axis=1), LABEL_POS: latent[:10, 4:], X_NOISE: data_init}
else:
feed_dict = {X_NOISE: data_init, LABEL_SIZE: latent_scale, LABEL_POS: latent_pos}
x_output = sess.run([x_final], feed_dict=feed_dict)[0]
if FLAGS.joint_shape:
im_name = "size_shape_combine_gentest.png"
else:
im_name = "size_scale_combine_gentest.png"
x_output_wrap = np.ones((10, 66, 66))
data_try_wrap = np.ones((10, 66, 66))
x_output_wrap[:, 1:-1, 1:-1] = x_output
data_try_wrap[:, 1:-1, 1:-1] = data_try
im_output = np.concatenate([x_output_wrap, data_try_wrap], axis=2).reshape(-1, 66*2)
impath = osp.join(save_exp_dir, im_name)
imsave(impath, im_output)
print("Successfully saved images at {}".format(impath))