def f_convert_sample_into_image(sample_z, eps_std):
with tf.variable_scope('model', reuse=True):
z = decoder(sample_z, eps_std=eps_std)
z = Z.unsqueeze2d(z, 2) # 8x8x12 -> 16x16x3
x = postprocess(z)
return x
def split2d_reverse(name, z, eps_std=None):
with tf.variable_scope(name):
z1 = Z.unsqueeze2d(z)
pz = split2d_prior(z1)
z2 = pz.sample
if eps_std is not None:
z2 = pz.sample2(pz.eps * tf.reshape(eps_std, [-1, 1, 1, 1]))
z = tf.concat([z1, z2], 3)
return z
def f_sample(y_A, y_B, eps_std):
with tf.variable_scope('model_A', reuse=True):
y_onehot_A = tf.cast(tf.one_hot(y_A, hps.n_y, 1, 0), 'float32')
_, sample, _ = prior("prior", y_onehot_A, hps)
z = sample(eps_std=eps_std)
z = decoder_A(z, eps_std=eps_std)
z = Z.unsqueeze2d(z, 2) # 8x8x12 -> 16x16x3
x_A = postprocess(z)
with tf.variable_scope('model_B', reuse=True):
y_onehot_B = tf.cast(tf.one_hot(y_B, hps.n_y, 1, 0), 'float32')
_, sample, _ = prior("prior", y_onehot_B, hps)
z = sample(eps_std=eps_std)
z = decoder_B(z, eps_std=eps_std)
z = Z.unsqueeze2d(z, 2) # 8x8x12 -> 16x16x3
x_B = postprocess(z)
return x_A, x_B
def f_sample(y, eps_std):
with tf.variable_scope('model', reuse=True):
y_onehot = tf.cast(tf.one_hot(y, hps.n_y, 1, 0), 'float32')
_, sample, _ = prior("prior", y_onehot, hps)
z = sample(eps_std=eps_std)
z, logdet_out = decoder(z, eps_std=eps_std)
z = Z.unsqueeze2d(z, 2) # 8x8x12 -> 16x16x3
x = postprocess(z)
return x, logdet_out
def f_decode(y, eps, reuse=True):
with tf.compat.v1.variable_scope('model', reuse=reuse):
y_onehot = tf.cast(tf.one_hot(y, hps.n_y, 1, 0), 'float32')
_, sample, _ = prior("prior", y_onehot, hps)
z = sample(eps=eps[-1])
z = decoder(z, eps=eps[:-1])
z = Z.unsqueeze2d(z, 2) # 8x8x12 -> 16x16x3
x = postprocess(z)
return x
def f_decode(y, eps, reuse=True):
with tf.variable_scope('model', reuse=reuse):
y_onehot = tf.cast(tf.one_hot(y, hps.n_y, 1, 0), 'float32')
_, sample, _ = prior("prior", y_onehot, hps)
z = sample(eps=eps[-1])
z = decoder(z, eps=eps[:-1])
z = Z.unsqueeze2d(z, 2) # 8x8x12 -> 16x16x3
x = postprocess(z)
return x
def f_decode(y, eps, model_name, reuse=True):
assert model_name == 'model_A' or model_name == 'model_B'
decoder = decoder_A if model_name == 'model_A' else decoder_B
with tf.variable_scope(model_name, reuse=reuse):
y_onehot = tf.cast(tf.one_hot(y, hps.n_y, 1, 0), 'float32')
_, sample, _ = prior("prior", y_onehot, hps)
z = sample(eps=eps[-1])
z = decoder(z, eps=eps[:-1])
z = Z.unsqueeze2d(z, 2) # 8x8x12 -> 16x16x3
x = postprocess(z)
return x
def split2d_reverse(name, z, eps, eps_std):
with tf.variable_scope(name):
z1 = Z.unsqueeze2d(z)
pz = split2d_prior(z1)
if eps is not None:
# Already sampled eps
z2 = pz.sample2(eps)
elif eps_std is not None:
# Sample with given eps_std
z2 = pz.sample2(pz.eps * tf.reshape(eps_std, [-1, 1, 1, 1]))
else:
# Sample normally
z2 = pz.sample
z = tf.concat([z1, z2], 3)
return z
def split2d_reverse(name, z, eps, eps_std):
with tf.variable_scope(name):
z1 = Z.unsqueeze2d(z)
pz = split2d_prior(z1)
if eps is not None:
# Already sampled eps
z2 = pz.sample2(eps)
elif eps_std is not None:
# Sample with given eps_std
z2 = pz.sample2(pz.eps * tf.reshape(eps_std, [-1, 1, 1, 1]))
else:
# Sample normally
z2 = pz.sample
z = tf.concat([z1, z2], 3)
return z
def split2d_reverse(name, z, hps, eps, eps_std):
with tf.variable_scope(name):
z1 = Z.unsqueeze2d(z)
pz = split2d_prior(z1, hps)
if eps is not None:
# Already sampled eps
z2 = pz.sample2(eps)
elif eps_std is not None:
# Sample with given eps_std
z2 = pz.sample2(pz.eps * tf.reshape(eps_std, [-1, 1, 1, 1]))
else:
# Sample normally
z2 = pz.sample
z = tf.concat([z1, z2], 3)
# z = tf.Print(z, data=[
# tf.reduce_max(pz.logsd), tf.reduce_min(pz.logsd)],
# message='split2d_prior logsd max/min')
return z
def _f_loss(x_A, y_A, x_B, y_B, is_training, reuse=False, init=False):
with tf.variable_scope('model_A', reuse=reuse):
y_onehot_A = tf.cast(tf.one_hot(y_A, hps.n_y, 1, 0), 'float32')
# Discrete -> Continuous
objective_A = tf.zeros_like(x_A, dtype='float32')[:, 0, 0, 0]
z_A = preprocess(x_A)
z_A = z_A + tf.random_uniform(tf.shape(z_A), 0, 1./hps.n_bins)
objective_A += - np.log(hps.n_bins) * np.prod(Z.int_shape(z_A)[1:])
# Encode
z_A = Z.squeeze2d(z_A, 2) # > 16x16x12
z_A, objective_A, eps_A = encoder_A(z_A, objective_A)
# Prior
hps.top_shape = Z.int_shape(z_A)[1:]
logp_A, _, _eps_A = prior("prior", y_onehot_A, hps)
objective_A += logp_A(z_A)
# Note that we learn the top layer so need to process z
z_A = _eps_A(z_A)
eps_A.append(z_A)
# Loss of eps and flatten latent code from another model
eps_flatten_A = tf.concat(
[tf.contrib.layers.flatten(e) for e in eps_A], axis=-1)
with tf.variable_scope('model_B', reuse=reuse):
y_onehot_B = tf.cast(tf.one_hot(y_B, hps.n_y, 1, 0), 'float32')
# Discrete -> Continuous
objective_B = tf.zeros_like(x_B, dtype='float32')[:, 0, 0, 0]
z_B = preprocess(x_B)
z_B = z_B + tf.random_uniform(tf.shape(z_B), 0, 1./hps.n_bins)
objective_B += - np.log(hps.n_bins) * np.prod(Z.int_shape(z_B)[1:])
# Encode
z_B = Z.squeeze2d(z_B, 2) # > 16x16x12
z_B, objective_B, eps_B = encoder_B(z_B, objective_B)
# Prior
hps.top_shape = Z.int_shape(z_B)[1:]
logp_B, _, _eps_B = prior("prior", y_onehot_B, hps)
objective_B += logp_B(z_B)
# Note that we learn the top layer so need to process z
z_B = _eps_B(z_B)
eps_B.append(z_B)
# Loss of eps and flatten latent code from another model
eps_flatten_B = tf.concat(
[tf.contrib.layers.flatten(e) for e in eps_B], axis=-1)
code_loss = 0.0
code_shapes = [[16, 16, 6], [8, 8, 12], [4, 4, 48]]
if hps.code_loss_type == 'B_all':
if not init:
""" Decode the code from another model and compute L2 loss
at pixel level
"""
def unflatten_code(fcode, code_shapes):
index = 0
code = []
bs = tf.shape(fcode)[0]
# bs = hps.local_batch_train
for shape in code_shapes:
code.append(tf.reshape(fcode[:, index:index+np.prod(shape)],
tf.convert_to_tensor([bs] + shape)))
index += np.prod(shape)
return code
code_others = unflatten_code(eps_flatten_A, code_shapes)
# code_others[-1] is z, and code_others[:-1] is eps
with tf.variable_scope('model_B', reuse=True):
_, sample, _ = prior("prior", y_onehot_B, hps)
code_last_others = sample(eps=code_others[-1])
code_decoded_others = decoder_B(
code_last_others, code_others[:-1])
code_decoded = Z.unsqueeze2d(code_decoded_others, 2)
x_B_recon = postprocess(code_decoded)
x_B_scaled = 1/255.0 * tf.cast(x_B, tf.float32)
x_B_recon_scaled = 1/255.0 * tf.cast(x_B_recon, tf.float32)
if hps.code_loss_fn == 'l1':
code_loss = tf.reduce_mean(tf.losses.absolute_difference(
x_B_scaled, x_B_recon_scaled))
elif hps.code_loss_fn == 'l2':
code_loss = tf.reduce_mean(tf.squared_difference(
x_B_scaled, x_B_recon_scaled))
else:
raise NotImplementedError()
elif hps.code_loss_type == 'code_all':
code_loss = tf.reduce_mean(
tf.squared_difference(eps_flatten_A, eps_flatten_B))
elif hps.code_loss_type == 'code_last':
dim = np.prod(code_shapes[-1])
code_loss = tf.reduce_mean(tf.squared_difference(
eps_flatten_A[:, -dim:], eps_flatten_B[:, -dim:]))
else:
raise NotImplementedError()
with tf.variable_scope('model_A', reuse=True):
# Generative loss
nobj_A = - objective_A
bits_x_A = nobj_A / (np.log(2.) * int(x_A.get_shape()[1]) * int(
x_A.get_shape()[2]) * int(x_A.get_shape()[3])) # bits per subpixel
bits_y_A = tf.zeros_like(bits_x_A)
classification_error_A = tf.ones_like(bits_x_A)
with tf.variable_scope('model_B', reuse=True):
# Generative loss
nobj_B = - objective_B
bits_x_B = nobj_B / (np.log(2.) * int(x_B.get_shape()[1]) * int(
x_B.get_shape()[2]) * int(x_B.get_shape()[3])) # bits per subpixel
bits_y_B = tf.zeros_like(bits_x_B)
classification_error_B = tf.ones_like(bits_x_B)
return (bits_x_A, bits_y_A, classification_error_A, eps_flatten_A,
bits_x_B, bits_y_B, classification_error_B, eps_flatten_B, code_loss)
