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, 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)