New conditional encoder, and variational

Author: vanpelt

keras-autoencoder/conditional_vae.py renamed to keras-autoencoder/conditional_autoencoder.py

@@ -1,35 +1,34 @@
 import matplotlib
 matplotlib.use("Agg")
-from wandb.keras import WandbCallback
-import wandb
-from keras.utils.generic_utils import get_custom_objects
-from keras import backend as K
-from keras.callbacks import Callback
-from keras.models import Model, load_model
-from keras import metrics #mse, binary_crossentropy
-from keras import layers
-import keras
-import plotly.graph_objs as go
-import plotly.plotly as py
-from sklearn import manifold
-from sklearn.decomposition import PCA
-import matplotlib.pyplot as plt
-import pdb
-import os
-import subprocess
-import pandas as pd
-import numpy as np
-import cv2
 import sys
+import cv2
+import numpy as np
+import pandas as pd
+import subprocess
+import os
+import pdb
+import matplotlib.pyplot as plt
+from sklearn.decomposition import PCA
+from sklearn import manifold
+import plotly.plotly as py
+import plotly.graph_objs as go
+import keras
+from keras import layers
+from keras.models import Model, load_model
+from keras.callbacks import Callback
+from keras import backend as K
+from keras.utils.generic_utils import get_custom_objects
+import wandb
+from wandb.keras import WandbCallback
 
 
-wandb.init(project="cvae")
+wandb.init()
 wandb.config.latent_dim = 2
-wandb.config.labels = [str(i) for i in range(10)]# ["Happy", "Sad"]
+wandb.config.labels = [str(i) for i in range(10)] #["Happy", "Sad"]
 wandb.config.batch_size = 128
 wandb.config.epochs = 25
-wandb.config.variational = False
-wandb.config.conditional = False
+wandb.config.conditional = True
+wandb.config.latent_vis = False
 wandb.config.dataset = "mnist"
 
 EMOTIONS = ["Angry", "Disgust", "Fear", "Happy", "Sad", "Surprise", "Neutral"]
@@ -66,24 +65,6 @@ def load_fer2013(filter_emotions=[]):
     return train_faces, train_emotions, val_faces, val_emotions
 
 
-def sample(args):
-    '''
-    Draws samples from a standard normal and scales the samples with
-    standard deviation of the variational distribution and shifts them
-    by the mean.
-
-    Args:
-        args: sufficient statistics of the variational distribution.
-
-    Returns:
-        Samples from the variational distribution.
-    '''
-    t_mean, t_log_var = args
-    t_sigma = K.sqrt(K.exp(t_log_var))
-    epsilon = K.random_normal(shape=K.shape(t_mean), mean=0., stddev=1.)
-    return t_mean + t_sigma * epsilon
-
-
 def concat_label(args):
     '''
     Converts 2d labels to 3d, i.e. [[0,1]] = [[[0,0],[0,0]],[[1,1],[1,1]]]
@@ -103,43 +84,46 @@ class ShowImages(Callback):
     '''
     Keras callback for logging predictions and a scatter plot of the latent dimension
     '''
+
     def on_epoch_end(self, epoch, logs):
         indicies = np.random.randint(X_test.shape[0], size=36)
         latent_idx = np.random.randint(X_test.shape[0], size=500)
         inputs = X_test[indicies]
         t_inputs = X_train[indicies]
         r_labels = y_test[indicies]
         rand_labels = np.random.randint(len(wandb.config.labels), size=35)
-        rand_labels = np.append(rand_labels, [len(wandb.config.labels) - 1]) # always add max label
+        # always add max label
+        rand_labels = np.append(rand_labels, [len(wandb.config.labels) - 1])
         labels = keras.utils.to_categorical(rand_labels)
         t_labels = y_train[indicies]
 
-        results = vae.predict([inputs, r_labels, labels])
-        t_results = vae.predict([t_inputs, t_labels, t_labels])
-        if wandb.config.variational:
-            output = encoder.predict([X_test[latent_idx], y_test[latent_idx]])
-            latent = output[0] #K.eval(sample(output))
-        else:
-            latent = encoder.predict([X_test[latent_idx], y_test[latent_idx]])
+        results = cae.predict([inputs, r_labels, labels])
+        t_results = cae.predict([t_inputs, t_labels, t_labels])
+        print("Max pixel value", t_results.max())
+        latent = encoder.predict([X_test[latent_idx], y_test[latent_idx]])
         # Plot latent space
-        # tsne = manifold.TSNE(n_components=2, init='pca', random_state=0)
-        latent_vis = PCA(n_components=2)
-        X = latent_vis.fit_transform(latent)
-        trace = go.Scatter(x=list(X[:, 0]), y=list(X[:, 1]),
-                           mode='markers', showlegend=False,
-                           marker=dict(color=list(np.argmax(y_test[latent_idx], axis=1)),
-                                       colorscale='Viridis',
-                                       size=8,
-                                       showscale=True))
-        fig = go.Figure(data=[trace])
-        wandb.log({"latent_vis": fig}, commit=False)
+        if wandb.config.latent_vis:
+            if wandb.config.latent_dim > 2:
+                # latent_vis = manifold.TSNE(n_components=2, init='pca', random_state=0)
+                latent_vis = PCA(n_components=2)
+                X = latent_vis.fit_transform(latent)
+            else:
+                X = latent
+            trace = go.Scatter(x=list(X[:, 0]), y=list(X[:, 1]),
+                            mode='markers', showlegend=False,
+                            marker=dict(color=list(np.argmax(y_test[latent_idx], axis=1)),
+                                        colorscale='Viridis',
+                                        size=8,
+                                        showscale=True))
+            fig = go.Figure(data=[trace])
+            wandb.log({"latent_vis": fig}, commit=False)
         # Always log training images
         wandb.log({
             "train_images": [wandb.Image(
                 np.hstack([t_inputs[i], res])) for i, res in enumerate(t_results)
             ]
         }, commit=False)
-        
+
         # Log image conversion when conditional
         if wandb.config.conditional:
             wandb.log({
@@ -153,27 +137,16 @@ def on_epoch_end(self, epoch, logs):
 def create_encoder(input_shape):
     '''
     Create an encoder with an optional class append to the channel.
-    Optionally outputting mean and stddev for a variational encoder
     '''
     encoder_input = layers.Input(shape=input_shape)
     label_input = layers.Input(shape=(len(wandb.config.labels),))
+    x = layers.Flatten()(encoder_input)
     if wandb.config.conditional:
-        x = layers.Lambda(concat_label, name="c")([encoder_input, label_input])
-    else:
-        x = encoder_input
-    x = layers.Conv2D(32, 3, padding='same',
-                      activation='relu', strides=(2, 2))(x)
-    x = layers.Conv2D(64, 3, padding='same',
-                      activation='relu', strides=(2, 2))(x)
-    x = layers.Flatten()(x)
-    x = layers.Dense(32, activation='relu')(x)
-    if wandb.config.variational:
-        t_mean = layers.Dense(wandb.config.latent_dim, name="latent_mean")(x)
-        t_log_var = layers.Dense(
-            wandb.config.latent_dim, name="latent_variance")(x)
-        output = [t_mean, t_log_var]
-    else:
-        output = layers.Dense(wandb.config.latent_dim, activation="relu")(x)
+        #x = layers.Lambda(concat_label, name="c")([encoder_input, label_input])
+        x = layers.concatenate([x, label_input], axis=-1)
+
+    x = layers.Dense(512, activation="relu")(x)
+    output = layers.Dense(wandb.config.latent_dim, activation="relu")(x)
 
     return Model([encoder_input, label_input], output, name='encoder')
 
@@ -185,15 +158,12 @@ def create_categorical_decoder():
     decoder_input = layers.Input(shape=(wandb.config.latent_dim,))
     label_input = layers.Input(shape=(len(wandb.config.labels),))
     if wandb.config.conditional:
-        x = layers.concatenate([decoder_input, label_input], axis=1)
+        x = layers.concatenate([decoder_input, label_input], axis=-1)
     else:
         x = decoder_input
-    x = layers.Dense(img_size // 2 * img_size // 2 * 32, activation='relu')(x)
-    x = layers.Reshape((img_size // 2, img_size // 2, 32))(x)
-    x = layers.Conv2D(32, 3, padding='same', activation='relu')(x)
-    x = layers.Conv2DTranspose(
-        32, 3, padding='same', activation='relu', strides=(2, 2))(x)
-    x = layers.Conv2D(1, 3, padding='same', activation='sigmoid')(x)
+    x = layers.Dense(512, activation='relu')(x)
+    x = layers.Dense(img_size * img_size, activation='sigmoid')(x)
+    x = layers.Reshape((img_size, img_size, 1))(x)
 
     return Model([decoder_input, label_input], x, name='decoder')
 
@@ -232,45 +202,21 @@ def create_categorical_decoder():
 
 encoder = create_encoder(input_shape=(img_size, img_size, 1))
 decoder = create_categorical_decoder()
-sampler = layers.Lambda(sample, name='sampler')
 
 image = layers.Input(shape=(img_size, img_size, 1))
 true_label = layers.Input(shape=(len(wandb.config.labels),))
 dest_label = layers.Input(shape=(len(wandb.config.labels),))
 output = encoder([image, true_label])
-if wandb.config.variational:
-    t_mean, t_log_var = output
-    t = sampler(output)
-else:
-    t = output
-t_decoded = decoder([t, dest_label])
-
-def vae_loss(y_true, y_pred):
-    ''' Negative variational lower bound used as loss function for training the variational auto-encoder. '''
-    # Reconstruction loss
-    rc_loss = metrics.mse(K.flatten(image), K.flatten(t_decoded)) #binary_crossentropy
-    #rc_loss *= img_size * img_size
-
-    if wandb.config.variational:
-        # Regularization term (KL divergence)
-        kl_loss = -0.5 * K.sum(1 + t_log_var
-                               - K.square(t_mean)
-                               - K.exp(t_log_var), axis=-1)
-        # Average over mini-batch
-        return K.mean(rc_loss + kl_loss)
-    else:
-        return K.mean(rc_loss)
-
-get_custom_objects().update({"vae_loss": vae_loss})
+t_decoded = decoder([output, dest_label])
 
 
-vae = Model([image, true_label, dest_label], t_decoded, name='vae')
-vae.compile(optimizer='rmsprop', loss=vae_loss)
+cae = Model([image, true_label, dest_label], t_decoded, name='vae')
+cae.compile(optimizer='rmsprop', loss="mse")
 
 print(X_train.shape, y_train.shape, X_test.shape, y_test.shape)
 
 if __name__ == '__main__':
-    vae.fit([X_train, y_train, y_train], X_train, epochs=wandb.config.epochs,
+    cae.fit([X_train, y_train, y_train], X_train, epochs=wandb.config.epochs,
             shuffle=True, batch_size=wandb.config.batch_size, callbacks=[ShowImages(), WandbCallback()],
             validation_data=([X_test, y_test, y_test], X_test))
     encoder.save("encoder.h5")