CVAE(条件自编码) Condition GAN (条件GAN) 和 VAE-GAN模型之间的区别之CVAE

# 使用CVAE(条件自编码) 训练fashion-mnist数据集

[code]import os
import time
import tensorflow as tf
import numpy as np

from ops import *
from utils import *

class CVAE(object):
model_name = "CVAE"     # name for checkpoint

def __init__(self, sess, epoch, batch_size, z_dim, dataset_name, checkpoint_dir, result_dir, log_dir):
self.sess = sess
self.dataset_name = dataset_name
self.checkpoint_dir = checkpoint_dir
self.result_dir = result_dir
self.log_dir = log_dir
self.epoch = epoch
self.batch_size = batch_size
self.mean = 0
self.var =1

if dataset_name == 'mnist' or dataset_name == 'fashion-mnist':
# parameters
self.input_height = 28
self.input_width = 28
self.output_height = 28
self.output_width = 28

self.z_dim = z_dim         # dimension of noise-vector
self.y_dim = 10         # dimension of condition-vector (label)
self.c_dim = 1

# train
self.learning_rate = 0.0002
self.beta1 = 0.5

# test
self.sample_num = 64  # number of generated images to be saved

# load mnist
self.data_X, self.data_y = load_mnist(self.dataset_name)

# get number of batches for a single epoch
self.num_batches = len(self.data_X) // self.batch_size
else:
print("********there is no other dataset to do *********")
raise NotImplementedError

# 编码器中输入的是真实图像和噪音向量
def encoder(self, x, y, is_training=True, reuse=False):
with tf.variable_scope("encoder", reuse=reuse):
y = tf.reshape(y, [self.batch_size, 1, 1, self.y_dim])
x = conv_cond_concat(x, y)

net = lrelu(conv2d(x, 64, 4, 4, 2, 2, name='en_conv1'))
net = lrelu(bn(conv2d(net, 128, 4, 4, 2, 2, name='en_conv2'), is_training=is_training, scope='en_bn2'))
net = tf.reshape(net, [self.batch_size, -1])
net = lrelu(bn(linear(net, 1024, scope='en_fc3'), is_training=is_training, scope='en_bn3'))
gaussian_params = linear(net, 2 * self.z_dim, scope='en_fc4')

mean = gaussian_params[:, :self.z_dim]
stddev = tf.nn.softplus(gaussian_params[:, self.z_dim:])

return mean, stddev

# 定义解码器的相关操作
def decoder(self, z, y, is_training=True, reuse=False):
with tf.variable_scope("decoder", reuse=reuse):

z = concat([z, y], 1)
net = tf.nn.relu(bn(linear(z, 1024, scope='de_fc1'), is_training=is_training, scope='de_bn1'))
net = tf.nn.relu(bn(linear(net, 128 * 7 * 7, scope='de_fc2'), is_training=is_training, scope='de_bn2'))
net = tf.reshape(net, [self.batch_size, 7, 7, 128])
net = tf.nn.relu(
bn(deconv2d(net, [self.batch_size, 14, 14, 64], 4, 4, 2, 2, name='de_dc3'), is_training=is_training,
scope='de_bn3'))

out = tf.nn.sigmoid(deconv2d(net, [self.batch_size, 28, 28, 1], 4, 4, 2, 2, name='de_dc4'))

return out

def build_model(self):

image_dims = [self.input_height, self.input_width, self.c_dim]
bs = self.batch_size

self.inputs = tf.placeholder(tf.float32, [bs] + image_dims, name='real_images')
self.y = tf.placeholder(tf.float32, [bs, self.y_dim], name='y')
self.z = tf.placeholder(tf.float32, [bs, self.z_dim], name='z')

# 编码器中返回的数据是均值和方差 经过运算之后返回数据
mu, sigma = self.encoder(self.inputs, self.y,  is_training=True, reuse=False)
z = mu + sigma * tf.random_normal(tf.shape(mu), 0, 1, dtype=tf.float32)

# 解码器输出真实图像
self.out = self.decoder(z, self.y, is_training=True, reuse=False)

# 定义loss函数
marginal_likelihood = tf.reduce_sum(self.inputs * tf.log(self.out) + (1 - self.inputs) * tf.log(1 - self.out), [1, 2])
KL_divergence = 0.5 * tf.reduce_sum(tf.square(mu) + tf.square(sigma) - tf.log(1e-8 + tf.square(sigma)) - 1, [1])

self.neg_loglikelihood = -tf.reduce_mean(marginal_likelihood)
self.KL_divergence = tf.reduce_mean(KL_divergence)

# 这个损失函数不是很懂 生成结果好 就这样用吧
self.loss = self.neg_loglikelihood + self.KL_divergence

# 定义优化器
t_vars = tf.trainable_variables()
with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
self.optim = tf.train.AdamOptimizer(self.learning_rate*5, beta1=self.beta1) \
.minimize(self.loss, var_list=t_vars)

# is_training设置为false生成图像 但是不参与模型修改参数
self.fake_images = self.decoder(self.z, self.y, is_training=False, reuse=True)

self.merged_summary_op = tf.summary.merge_all()

def train(self):

tf.global_variables_initializer().run()

# 标签使用的是前batch_size个图像
self.sample_z = np.random.normal(self.mean, self.var, (self.batch_size, self.z_dim)).astype(np.float32)
self.test_labels = self.data_y[0:self.batch_size]

start_epoch = 0
start_batch_id = 0
counter = 1

start_time = time.time()
for epoch in range(start_epoch, self.epoch):

for idx in range(start_batch_id, self.num_batches):
batch_images = self.data_X[idx*self.batch_size:(idx+1)*self.batch_size]
batch_labels = self.data_y[idx * self.batch_size:(idx + 1) * self.batch_size]
batch_z = np.random.uniform(-1, 1, [self.batch_size, self.z_dim]).astype(np.float32)

_, summary_str, loss, nll_loss, kl_loss = self.sess.run([self.optim, self.merged_summary_op, self.loss, self.neg_loglikelihood, self.KL_divergence],
feed_dict={self.inputs: batch_images, self.y: batch_labels, self.z: batch_z})

counter += 1
print("Epoch: [%2d] [%4d/%4d] time: %4.4f, loss: %.8f, nll: %.8f, kl: %.8f" \
% (epoch, idx, self.num_batches, time.time() - start_time, loss, nll_loss, kl_loss))

# save training results for every 300 steps
if np.mod(counter, 300) == 0:
samples = self.sess.run(self.fake_images,
feed_dict={self.z: self.sample_z, self.y: self.test_labels})
tot_num_samples = min(self.sample_num, self.batch_size)
manifold_h = int(np.floor(np.sqrt(tot_num_samples)))
manifold_w = int(np.floor(np.sqrt(tot_num_samples)))
save_images(samples[:manifold_h * manifold_w, :, :, :], [manifold_h, manifold_w],
'./' + check_folder(self.result_dir + '/' + self.model_dir) + '/' + self.model_name + '_train_{:02d}_{:04d}.png'.format(
epoch, idx))
start_batch_id = 0

@property
def model_dir(self):
return "{}_{}_{}_{}".format(
self.model_name, self.dataset_name,
self.batch_size, self.z_dim)

训练的结果：