GDPP-VAE Stacked MNIST

gdppvae_stackedMnist.py

@@ -0,0 +1,385 @@
+import os, random, numpy as np
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = "3"
+random_seed = 100
+random.seed(random_seed)
+np.random.seed(random_seed)
+
+from ops import *
+from scipy.stats import entropy
+from tflib.linear import Linear
+from tflib.batchnorm import Batchnorm
+from tflib.deconv2d import Deconv2D
+from tflib.conv2d import Conv2D
+import matplotlib.gridspec as gridspec
+import matplotlib.pyplot as plt
+# from tensorflow.examples.tutorials.mnist import input_data, mnist
+
+dpp_weight = 100.0	# Used to scale the DPP-measure to match the latent&generator loss ranges.
+num_epochs = 10
+num_eval_samples = 26000
+
+
+def compute_diversity_loss(h_fake, h_real):
+    def compute_diversity(h):
+        h = tf.nn.l2_normalize(h, 1)
+        Ly = tf.tensordot(h, tf.transpose(h), 1)
+        eig_val, eig_vec = tf.self_adjoint_eig(Ly)
+        return eig_val, eig_vec
+
+    def normalize_min_max(eig_val):
+        return tf.div(tf.subtract(eig_val, tf.reduce_min(eig_val)),
+                      tf.subtract(tf.reduce_max(eig_val), tf.reduce_min(eig_val)))  # Min-max-Normalize Eig-Values
+
+    fake_eig_val, fake_eig_vec = compute_diversity(h_fake)
+    real_eig_val, real_eig_vec = compute_diversity(h_real)
+    # Used a weighing factor to make the two losses operating in comparable ranges.
+    eigen_values_loss = 0.0001 * tf.losses.mean_squared_error(labels=real_eig_val, predictions=fake_eig_val)
+    eigen_vectors_loss = -tf.reduce_sum(tf.multiply(fake_eig_vec, real_eig_vec), 0)
+    normalized_real_eig_val = normalize_min_max(real_eig_val)
+    weighted_eigen_vectors_loss = tf.reduce_sum(tf.multiply(normalized_real_eig_val, eigen_vectors_loss))
+    return eigen_values_loss + weighted_eigen_vectors_loss
+
+
+def LeakyReLU(x, alpha=0.2):
+    return tf.maximum(alpha * x, x)
+
+
+class LatentAttention():
+    def __init__(self, dpp_weight, num_eval_samples=26000):
+        self.n_samples = 50000
+        self.dpp_weight = dpp_weight
+        self.num_eval_samples = num_eval_samples
+        self.n_hidden = 500
+        self.n_z = 128
+        self.im_size = 28
+        self.dim_h = 64
+        self.batchsize = 100
+        self.dim_x = 3 * self.im_size * self.im_size
+
+        self.images = tf.placeholder(tf.float32, [None, 3 * 784])
+        image_matrix = tf.reshape(self.images, [-1, 28, 28, 3])
+        z_mean, z_stddev, h_real = self.recognition(image_matrix)
+        samples = tf.random_normal([self.batchsize, self.n_z], 0, 1, dtype=tf.float32)
+        guessed_z = z_mean + (z_stddev * samples)
+
+        self.generated_images = self.generation(guessed_z)
+        generated_flat = tf.reshape(self.generated_images, [self.batchsize, 3 * 28 * 28])
+
+        self.dpp_loss = tf.constant(0.0)
+        if self.dpp_weight > 0:
+            samples_fake = tf.random_normal([self.batchsize, self.n_z], 0, 1, dtype=tf.float32)
+            generated_fake = self.generation(samples_fake, reuse=True)
+            _, _, h_fake = self.recognition(generated_fake, reuse=True)
+            self.dpp_loss = compute_diversity_loss(h_fake, h_real)
+
+        self.generation_loss = -tf.reduce_sum(
+            self.images * tf.log(1e-8 + generated_flat) + (1 - self.images) * tf.log(1e-8 + 1 - generated_flat), 1)
+        self.full_gen_loss = self.generation_loss + self.dpp_weight * self.dpp_loss
+
+        self.latent_loss = 0.5 * tf.reduce_sum(
+            tf.square(z_mean) + tf.square(z_stddev) - tf.log(tf.square(z_stddev)) - 1, 1)
+        self.cost = tf.reduce_mean(self.generation_loss + self.latent_loss)
+        self.optimizer = tf.train.AdamOptimizer(0.001).minimize(self.cost)
+
+        self.eval_samples = tf.random_normal([num_eval_samples, self.n_z], 0, 1, dtype=tf.float32)
+        self.eval_generated_fakes = self.generation(self.eval_samples, reuse=True)
+
+    # encoder/ discriminator
+    def recognition(self, input_images, reuse=False):
+        with tf.variable_scope("recognition") as scope:
+            if reuse:
+                scope.reuse_variables()
+
+            im = tf.reshape(input_images, [-1, 3, self.im_size, self.im_size])
+
+            conv1 = Conv2D('Discriminator.1', 3, self.dim_h, 5, im, stride=2)
+            out_conv1 = LeakyReLU(conv1)
+
+            conv2 = Conv2D('Discriminator.2', self.dim_h, 2 * self.dim_h, 5, out_conv1, stride=2)
+            out_conv2 = LeakyReLU(conv2)
+
+            conv3 = Conv2D('Discriminator.3', 2 * self.dim_h, 4 * self.dim_h, 5, out_conv2, stride=2)
+            out_conv3 = LeakyReLU(conv3)
+
+            # dim_h = 64 |||| w=4, h=4, num_channel=4*64
+            h2_flat = tf.reshape(out_conv3, [-1, 4 * 4 * 4 * self.dim_h])
+            w_mean = Linear('Discriminator.w_mean', 4 * 4 * 4 * self.dim_h, self.n_z, h2_flat)
+            w_stddev = Linear('Discriminator.w_stddev', 4 * 4 * 4 * self.dim_h, self.n_z, h2_flat)
+
+        return w_mean, w_stddev, h2_flat
+
+    # decoder/ generator
+    def generation(self, z, reuse=False):
+        with tf.variable_scope("generation") as scope:
+            if reuse:
+                scope.reuse_variables()
+            fc1 = Linear('Generator.Input', self.n_z, 4 * 4 * 4 * self.dim_h, z)
+            fc1 = Batchnorm('Generator.BN1', [0], fc1)
+            fc1 = tf.nn.relu(fc1)
+            out_fc1 = tf.reshape(fc1, [-1, 4 * self.dim_h, 4, 4])
+
+            deconv1 = Deconv2D('Generator.2', 4 * self.dim_h, 2 * self.dim_h, 5, out_fc1)
+            deconv1 = Batchnorm('Generator.BN2', [0, 2, 3], deconv1)
+            deconv1 = tf.nn.relu(deconv1)
+            out_deconv1 = deconv1[:, :, :7, :7]
+
+            deconv2 = Deconv2D('Generator.3', 2 * self.dim_h, self.dim_h, 5, out_deconv1)
+            deconv2 = Batchnorm('Generator.BN3', [0, 2, 3], deconv2)
+            out_deconv2 = tf.nn.relu(deconv2)
+
+            deconv3 = Deconv2D('Generator.5', self.dim_h, 3, 5, out_deconv2)
+            out_deconv3 = tf.sigmoid(deconv3)
+
+            return tf.reshape(out_deconv3, [-1, self.dim_x])
+
+    def inf_train_gen(self, dataset_path, BATCH_SIZE):
+        ds = np.load(dataset_path).item()
+        images, labels = ds['images'], ds['labels']
+        ds_size = labels.shape[0]
+        while True:
+            for i in range(int(ds_size / BATCH_SIZE)):
+                start = i * BATCH_SIZE
+                end = (i + 1) * BATCH_SIZE
+                yield images[start:end], labels[start:end]
+
+    def save_fig_color(self, samples, out_path, idx):
+        fig = plt.figure(figsize=(5, 5))
+        gs = gridspec.GridSpec(5, 5)
+        gs.update(wspace=0.05, hspace=0.05)
+
+        for i, sample in enumerate(samples):
+            ax = plt.subplot(gs[i])
+            plt.axis('off')
+            ax.set_xticklabels([])
+            ax.set_yticklabels([])
+            ax.set_aspect('equal')
+            s = sample.reshape(3, 28, 28)
+            s = s.transpose(1, 2, 0)
+            plt.imshow(s, cmap='Greys_r')
+
+        plt.savefig(out_path + '/{}.png'.format(str(idx).zfill(3)), bbox_inches='tight')
+        plt.close(fig)
+
+    def train(self, data_dir, num_epochs=10):
+        gen = self.inf_train_gen(data_dir + 'stacked_train.npy', self.batchsize)
+        visualization, _ = next(gen)
+        run_config = tf.ConfigProto()
+        run_config.gpu_options.allow_growth = True
+        with tf.Session(config=run_config) as sess:
+            np.random.seed(random_seed)
+            tf.set_random_seed(random_seed)
+            sess.run(tf.global_variables_initializer())
+            for epoch in range(num_epochs):
+                for idx in range(int(self.n_samples / self.batchsize)):
+                    batch, _ = next(gen)
+                    _, dpp_loss, gen_loss, lat_loss = sess.run(
+                        (self.optimizer, self.dpp_loss, self.generation_loss, self.latent_loss),
+                        feed_dict={self.images: batch})
+                    if idx % (self.n_samples - 3) == 0:
+                        print "epoch %d: genloss %f latloss %f dpploss %f" % (
+                        epoch, np.mean(gen_loss), np.mean(lat_loss), np.mean(dpp_loss))
+                        generated_test = sess.run(self.generated_images, feed_dict={self.images: visualization})
+                        generated_test = generated_test.reshape(self.batchsize, 28 * 28 * 3)
+                        self.save_fig_color(generated_test[:25], data_dir.replace('/data', '/Results'), epoch)
+
+            np.random.seed(random_seed)
+            tf.set_random_seed(random_seed)
+            eval_samples = np.array(sess.run(self.eval_generated_fakes)).reshape(self.num_eval_samples, 3,
+                                                                                 self.im_size ** 2)
+            return eval_samples
+
+
+####################################################################################
+# ------------------------------ Evaluation Network ------------------------------ #
+####################################################################################
+class MNIST(object):
+    def __init__(self, model_name):
+        # Create the model
+        self.model_name = model_name
+        self.x = tf.placeholder(tf.float32, [None, 784])
+        # Define loss and optimizer
+        self.y_ = tf.placeholder(tf.float32, [None, 10])
+        self.keep_prob = tf.placeholder(tf.float32)
+        self.W_conv1 = weight_variable([5, 5, 1, 32])
+        self.b_conv1 = bias_variable([32])
+
+        self.W_conv2 = weight_variable([5, 5, 32, 64])
+        self.b_conv2 = bias_variable([64])
+
+        self.W_fc1 = weight_variable([7 * 7 * 64, 1024])
+        self.b_fc1 = bias_variable([1024])
+
+        self.W_fc2 = weight_variable([1024, 10])
+        self.b_fc2 = bias_variable([10])
+
+    def conv2d(self, x, W):
+        """conv2d returns a 2d convolution layer with full stride."""
+        return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')
+
+    def deepnn(self, x):
+        x_image = tf.reshape(x, [-1, 28, 28, 1])
+        h_conv1 = tf.nn.relu(self.conv2d(x_image, self.W_conv1) + self.b_conv1)
+        h_pool1 = max_pool_2x2(h_conv1)
+        h_conv2 = tf.nn.relu(self.conv2d(h_pool1, self.W_conv2) + self.b_conv2)
+        h_pool2 = max_pool_2x2(h_conv2)
+        h_pool2_flat = tf.reshape(h_pool2, [-1, 7 * 7 * 64])
+        h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, self.W_fc1) + self.b_fc1)
+        h_fc1_drop = tf.nn.dropout(h_fc1, self.keep_prob)
+        y_conv = tf.matmul(h_fc1_drop, self.W_fc2) + self.b_fc2
+        return y_conv
+
+    def build_model(self, mnist):
+        # Build the graph for the deep net
+        y_conv = self.deepnn(self.x)
+
+        cross_entropy = tf.nn.softmax_cross_entropy_with_logits(labels=self.y_,
+                                                                logits=y_conv)
+        cross_entropy = tf.reduce_mean(cross_entropy)
+
+        params = tf.trainable_variables()
+        train_step = tf.train.AdamOptimizer(
+            1e-4).minimize(cross_entropy, var_list=params)
+
+        correct_prediction = tf.equal(
+            tf.argmax(y_conv, 1), tf.argmax(self.y_, 1))
+        correct_prediction = tf.cast(correct_prediction, tf.float32)
+        accuracy = tf.reduce_mean(correct_prediction)
+
+        with tf.Session() as sess:
+            sess.run(tf.global_variables_initializer())
+            saver = tf.train.Saver()
+            for i in xrange(20000):
+                batch = mnist.train.next_batch(50)
+                if i % 100 == 0:
+                    train_accuracy = accuracy.eval(
+                        feed_dict={self.x: batch[0], self.y_: batch[1], self.keep_prob: 1.0})
+                train_step.run(
+                    feed_dict={self.x: batch[0], self.y_: batch[1], self.keep_prob: 0.5})
+
+            saver.save(sess, self.model_name)
+
+    def classify(self, inputs):
+        y_conv = self.deepnn(self.x)
+        sess = tf.Session()
+        saver = tf.train.Saver()
+
+        saver.restore(sess, self.model_name)
+        # from IPython import embed; embed()
+        y_pred = sess.run(y_conv, feed_dict={
+            self.x: inputs, self.keep_prob: 1.0})
+        y = np.argmax(y_pred, 1)
+        return y
+
+
+##################################################################################
+# ------------------------------ Helper Functions ------------------------------ #
+##################################################################################
+
+def LeakyReLU(x, alpha=0.2):
+    return tf.maximum(alpha * x, x)
+
+
+class batch_norm(object):
+    def __init__(self, epsilon=1e-5, momentum=0.9, name="batch_norm"):
+        with tf.variable_scope(name):
+            self.epsilon = epsilon
+            self.momentum = momentum
+            self.name = name
+
+    def __call__(self, x, train=True):
+        return tf.contrib.layers.batch_norm(x,
+                                            decay=self.momentum,
+                                            updates_collections=None,
+                                            epsilon=self.epsilon,
+                                            scale=True,
+                                            is_training=train,
+                                            scope=self.name)
+
+
+def create_dir(dir_path):
+    if not os.path.exists(dir_path):
+        os.makedirs(dir_path)
+    return dir_path
+
+
+def max_pool_2x2(x):
+    """max_pool_2x2 downsamples a feature map by 2X."""
+    return tf.nn.max_pool(x, ksize=[1, 2, 2, 1],
+                          strides=[1, 2, 2, 1], padding='SAME')
+
+
+def weight_variable(shape):
+    """weight_variable generates a weight variable of a given shape."""
+    initial = tf.truncated_normal(shape, stddev=0.1)
+    return tf.Variable(initial)
+
+
+def bias_variable(shape):
+    """bias_variable generates a bias variable of a given shape."""
+    initial = tf.constant(0.1, shape=shape)
+    return tf.Variable(initial)
+
+
+def get_dist(p, n):
+    pk = np.zeros(n)
+    for x in p:
+        pk[x] += 1
+    for i in range(n):
+        pk[i] = pk[i] * 1.0 / len(p)
+    return pk
+
+
+def inf_train_gen(dataset_path, batch_size=64):
+    ds = np.load(dataset_path).item()
+    images, labels = ds['images'], ds['labels']
+    ds_size = labels.shape[0]
+    while True:
+        for i in range(int(ds_size / batch_size)):
+            start = i * batch_size
+            end = (i + 1) * batch_size
+            yield images[start:end], labels[start:end]
+
+
+###############################################################################
+# ------------------------------ Main Function ------------------------------ #
+###############################################################################
+
+if __name__ == '__main__':
+    main_dir = '/home/elfeki/Workspace/GDPP-VAE/MNIST/'
+    data_dir = main_dir + '/data/'
+
+    # # To generate the Stacked-MNIST data- One time
+    # from generate_stackedMnist_data import generate_stacked_mnist
+    # create_dir(data_dir)
+    # generate_stacked_mnist(data_dir)
+
+    model = LatentAttention(dpp_weight, num_eval_samples)
+    samples = model.train(data_dir, num_epochs)
+
+    np.random.seed(1)
+    tf.reset_default_graph()
+    tf.set_random_seed(random_seed)
+    nnet = MNIST(main_dir)
+
+    # # To train the CNN-classifier for evaluating the model- One time
+    # mnist = input_data.read_data_sets(data_dir, one_hot=True)
+    # nnet.build_model(mnist)
+
+    digit_1 = nnet.classify(samples[:, 0])
+    digit_2 = nnet.classify(samples[:, 1])
+    digit_3 = nnet.classify(samples[:, 2])
+    y_pred = []
+    for i in range(len(digit_1)):
+        y_pred.append(digit_1[i] * 100 + digit_2[i] * 10 + digit_3[i])
+    x = np.unique(y_pred)
+
+    ds = np.load(data_dir + 'stacked_train.npy').item()
+    labels = np.array(ds['labels'])
+    y_true = [np.argmax(lb) for lb in labels]
+    qk = get_dist(y_true, 1000)
+    pk = get_dist(y_pred, 1000)
+    kl_score = entropy(pk, qk)
+    print("#Modes: %d, KL-score: %.3f\n\n" % (len(x), kl_score))
+