{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Supervised Variational Auto-Encoder on MNIST\n\nThe goal of this notebook is to explore implementing a partially-labeled VAE on the MNIST data-set using Tensorflow Keras.","execution_count":null},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import math, re, os\nimport tensorflow as tf\nimport numpy as np\nimport pandas as pd\nfrom matplotlib import pyplot as plt\nfrom kaggle_datasets import KaggleDatasets\nfrom tqdm import tqdm\nprint( \"Tensorflow version \" + tf.__version__ )\nAUTO = tf.data.experimental.AUTOTUNE\nstrategy = tf.distribute.get_strategy() # default distribution strategy in Tensorflow. Works on CPU and single GPU.","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Configuration\n\nSet up the access to the labeled training dataset; since we want to be able to generate different digits on command, we only use the expert-labeled dataset to train the network.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"EPOCHS = 16\n\nGCS_PATH = KaggleDatasets().get_gcs_path()\n\nTRAINING_FILENAMES = tf.io.gfile.glob( GCS_PATH + '/train.csv' )\nCLASSES = [ str( i ) for i in range( 10 ) ] \nIMAGE_SIZE = ( 28, 28, 1 )                                                                                                  # 100 - 102","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Visualization Utilities\n\nThese functions just help display the data we generate throughout the notebook","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def displayImageBatch( images, labels=None, n=None ):\n    # choose to display\n    sample = ( np.random.choice( len( images ), n, replace=False )\n               if n else \n               np.arange( len( images ) ) )\n    \n    cols = 5\n    rows = ( len( sample ) - 1 ) // cols + 1\n    plt.figure( figsize=( 3 * cols, 3 * rows ) )\n    for i, s in enumerate( sample ):\n        plt.subplot( rows, cols, i + 1 )\n        plt.axis( 'off' )\n        plt.imshow( images[ s ], cmap='gray' )\n        if not labels is None:\n            plt.title( labels[ s ] )\n    plt.show()\n    \ndef displayTrainingCurve( training, title='Model Training' ):\n    plt.plot( training )\n    plt.title( title )\n    plt.xlabel( 'epoch' )\n    plt.ylabel( 'loss' )\n    plt.show()\n\ndef displayHist( data, title, nBins=10 ):\n    plt.hist( data, bins=nBins )\n    plt.title( title )\n    plt.xlabel( 'value' )\n    plt.ylabel( 'count' )\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Import Dataset\n\nPandas imports the raw data from csv, which we parse into useful formats:\n- X_train, with the normalized pixel data\n- label, with the class label for each sample\n- Y_train, with 1-hot encoding of class label for each sample","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# import training samples\nallTrain = pd.read_csv( '/kaggle/input/digit-recognizer/train.csv' )\nn = len( allTrain )\nX_train = allTrain.drop( 'label', 1 ).to_numpy( dtype=np.float32 ).reshape( ( -1, *IMAGE_SIZE ) ) / 255.\nlabel = allTrain[ 'label' ].to_numpy()\ndel allTrain\n\nclasses = np.unique( label )\nnClasses = len( classes )\nY_train = np.zeros( ( n, nClasses ), dtype=np.float32 )\nY_train[ np.arange( n ), label ] = 1.0\nprint( \"data import complete\" )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Visualize Original Dataset\n\nInspect the imported data to see what our network will be trying to reproduce; serves as a good sanity check on our imports.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"print( \"Training data shape:\", X_train.shape )\nprint( \"Training data label shape:\", Y_train.shape )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# run this cell again for next set of images\ndisplayImageBatch( X_train.reshape( -1, *IMAGE_SIZE[ :-1 ] ), label, 25 )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model\n\nConstruct the VAE itself using Keras Model along with directly chaining layers together. The architecture is built in two parts, encoder and decoder, then packaged into three Models:\n- The encoder Model contains only the portion which translates image samples into the latent space\n- The decoder Model contains only the portion which translates latent-space samples and labels into images\n- The trainer Model contains everything, and is used to train the VAE\nBy using layers in multiple models, we get all the nice Keras API for all the different things we want to do with these models. We do use a custom layer to produce the encoder loss to hack around some of the API's restrictions.\n\nThe nLatent can be adjusted here to change the size of the latent space.\n\nCautionary note: to adapt this model for use with other image sizes, changing IMAGE_SIZE is not sufficient because the decoder output size is dependent on the strides of its convolutional layers and its feed-in size. It could be automated with some restrictions, but that capability is elided for sake of more straight-forward code here.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"tfl = tf.keras.layers\nnLatent = 50\nwith strategy.scope():\n    \n    log2pi = tf.math.log( 2. * np.pi )\n    @tf.function\n    def log_normal_pdf( sample, mean, logvar, raxis=1 ):\n        return tf.reduce_sum(\n            -.5 * ( ( sample - mean ) ** 2. * tf.exp( -logvar ) + logvar + log2pi ), axis=raxis )\n    \n    @tf.function\n    def decoderLoss( y_true, y_pred ):\n            '''sum error instead of mean to defeat vanishing gradients in encoder'''\n            cross_ent = tf.nn.sigmoid_cross_entropy_with_logits( logits=y_pred, labels=y_true )\n            return tf.reduce_mean( cross_ent, axis=[ 1, 2, 3 ] )\n    \n    class EncoderLoss( tfl.Layer ):\n        def __init__( self, *args, **kwargs ):\n            super( EncoderLoss, self ).__init__( trainable=False, dynamic=True, *args, **kwargs )\n        \n        def compute_output_shape( self, input_shape ):\n            return type( input_shape )( ( None, ) )\n            \n        @tf.function\n        def call( self, inputs ):\n            mean, logvar = tf.split( inputs[ 0 ], num_or_size_splits=2, axis=1 )\n            logpz = log_normal_pdf( inputs[ 1 ], 0., 0. )\n            logqz_x = log_normal_pdf( inputs[ 1 ], mean, logvar )\n            \n            loss = logqz_x - logpz\n            reduce_axes = [ a for a in range( 1, len( loss.shape ) ) ]\n            if reduce_axes:\n                loss = tf.reduce_mean( loss, axis=reduce_axes )\n            return loss\n    \n    class Reparametrize( tfl.Layer ):\n        def __init__( self, *args, **kwargs ):\n            super( Reparametrize, self ).__init__( trainable=False, dynamic=True, *args, **kwargs )\n        \n        def compute_output_shape( self, input_shape ):\n            return type( input_shape )( [ ( s // 2 ) if i == 1 else s\n                                          for i, s in enumerate( input_shape ) ] )\n        \n        @tf.function\n        def call( self, inputs ):\n            mean, logvar = tf.split( inputs, num_or_size_splits=2, axis=1 )\n            noise = tf.random.normal( shape=mean.shape )\n            return noise * tf.exp( logvar * .5 ) + mean\n    \n    def buildVaeModels( latentDim=50, superDim=0 ):\n    \n        trainerInputs = []\n        trainerLosses = {}\n        \n        # encoder layers\n        encodeInput = tfl.Input( shape=IMAGE_SIZE )\n        encoderInputs = [ encodeInput ]\n        trainerInputs += encoderInputs\n        prev = encodeInput\n        for c in [ 32, 64 ]:\n            prev = tfl.Conv2D(\n                filters=c,\n                kernel_size=3,\n                strides=( 2, 2 ),\n                activation='relu' )( prev )\n        prev = tfl.Flatten()( prev )\n        for d in []:\n            prev = tfl.Dense( d, activation='relu' )( prev )\n        encodeOutput = tfl.Dense( 2 * latentDim, activation='relu' )( prev )\n        \n        # decoder layers\n        latentInput = tfl.Input( shape=( latentDim, ) )\n        decoderInputs = [ latentInput ]\n        prev = latentInput\n        if superDim > 0:\n            superInput = tfl.Input( shape=( superDim, ) )\n            decoderInputs += [ superInput ]\n            trainerInputs += [ superInput ]\n            prev = tf.keras.layers.concatenate( decoderInputs )\n        for d in [ ( 7 * 7 * 32 ) ]:\n            prev = tfl.Dense( d, activation='relu' )( prev )\n        prev = tfl.Reshape( target_shape=( 7, 7, 32 ) )( prev )\n        for c in [ 64, 32 ]:\n            prev = tfl.Conv2DTranspose( filters=c, kernel_size=3, strides= ( 2 , 2 ), padding=\"SAME\", activation='relu' )( prev )\n        decodeOutput = tfl.Conv2DTranspose( filters=1, kernel_size=3, strides= ( 1 , 1 ), padding=\"SAME\" )( prev )\n\n        # assemble models\n        encoder = tf.keras.Model( inputs=encoderInputs, outputs=[ encodeOutput ], name=\"Encoder\" )\n        decoder = tf.keras.Model( inputs=decoderInputs, outputs=[ decodeOutput ], name=\"Decoder\" )\n        \n        # trainer model\n        transitions = [ Reparametrize( name='Reparameter' )( encoder.layers[ -1 ].output ) ]\n        trainerLosses[ 'eLoss' ] = 'mae'\n        if superDim > 0:\n            transitions += trainerInputs[ -1: ]\n        trainerOutputs = [ EncoderLoss( name='eLoss')( [ encoder.output, transitions[ 0 ] ] ), decoder( transitions ) ]\n        trainerLosses[ 'Decoder' ] = decoderLoss\n        trainer = tf.keras.Model( inputs=trainerInputs,\n                                  outputs=trainerOutputs, name=\"Trainer\" )\n        trainer.compile( tf.keras.optimizers.Adam( 1e-4 ), loss=trainerLosses, loss_weights=[ 1, np.prod( IMAGE_SIZE ) ] )\n        \n        return trainer, encoder, decoder\n        \n    trainer, encoder, decoder = buildVaeModels( latentDim=nLatent, superDim=nClasses )\n    trainer.summary()\n    encoder.summary()\n    decoder.summary()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training\n\nUsing the all-inclusive trainer model, all parts of the VAE can be trained like any other Keras Model.\nNotice the array of zeros passed in with the targets, which is important for targeting the complicated layer-based loss used in the encoder.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"history = trainer.fit( [ X_train, Y_train ], [ np.zeros( len( X_train ) ), X_train ], epochs=EPOCHS, batch_size=10, verbose=1 )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"displayTrainingCurve( history.history[ 'loss' ], 'Weighted Total Loss' )\ndisplayTrainingCurve( history.history[ 'eLoss_loss' ], 'Encoder Loss' )\ndisplayTrainingCurve( history.history[ 'Decoder_loss' ], 'Decoder Loss' )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Distribution of Training Set in Latent Space\n\nExplore the performance of the trained encoder to see how well it maps to a standard normal distribution, examining both for the whole training set and the individual classes separately. Pay special attention to the covariance matrices, which also show how well the model utilizes the latent space.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# encode training set\nXE_train = encoder.predict( X_train )\nXE_mean = XE_train[ :, :XE_train.shape[ 1 ] // 2 ]\nXE_var = np.exp( XE_train[ :, XE_train.shape[ 1 ] // 2: ] )\n# split classes\nclass_mean = [ XE_mean[ label == l ] for l in range( nClasses ) ]\nclass_var = [ XE_var[ label == l ] for l in range( nClasses ) ]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# covariance matrix for distribution of means in latent space\ntotalCov = np.cov( XE_mean, rowvar=False )\n\nplt.figure( figsize=( 15, 15 ) )\nplt.title( 'Training Latent-Space Covariance' )\nplt.imshow( totalCov )\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# total partial distribution as mean of means and mean of variances (l2 norm of stddev) assuming independent gaussians\ntotalCenter = np.mean( XE_mean, axis=0 )\ntotalVar = np.mean( XE_var, axis=0 )\n\ndisplayHist( totalCenter, 'Encoder Latent Means' )\ndisplayHist( totalVar, ' Encoder Latent Variance ' )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# covariance matrix for distributions of means in latent space for each class\nclassCov = [ np.cov( m, rowvar=False ) for m in class_mean ]\nfor i, cov in enumerate( classCov ):\n    plt.figure( figsize=( 5, 5 ) )\n    plt.title( 'Training Latent-Space Covariance for %d' % i )\n    plt.imshow( cov )\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"# partial distribution as mean of means and mean of variances (l2 norm of stddev) assuming independent gaussians for each class\nclassCenter = [ np.mean( class_mean[ i ], axis=0 ) for i in range( nClasses ) ]\nclassVar = [ np.mean( class_var[ i ], axis=0 ) for i in range( nClasses ) ]\n\nfor i, cc, cv in zip( range( nClasses ), classCenter, classVar ):\n    displayHist( cc, 'Class %d Encoder Latent Means' % i )\n    displayHist( cv, 'Class %d Encoder Latent Variance ' % i )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Reconstruction for Training Set\n\nVisually inspect the performance of the model as an auto-encoder to sanity check the training results.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"XR_train = tf.sigmoid( decoder.predict( [ XE_mean, Y_train ] ) ).numpy()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# run this cell again for next set of images\ndisplayImageBatch( XR_train.reshape( -1, *IMAGE_SIZE[ :-1 ] ), label, 25 )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Generate New Samples\n\nFinally, the goal of the exercise: use the decoder to dream new instances.\n\nWe demonstrate the result of dreaming a sample with the standard normal as well as with correction for the observed distributions of training samples in latent spece. Note how little difference is made between the distributions.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# random sample setup; run this to generate new samples\nstddev = 1.\nnSamples = 25\nstd_sample = tf.random.normal( shape=( nSamples, nLatent ), stddev=stddev )\nlabel_sample = np.random.choice( nClasses, nSamples )\n\nlabel_sample_vec = np.zeros( ( nSamples, nClasses ), dtype=np.float32 )\nlabel_sample_vec[ np.arange( nSamples ), label_sample ] = 1.0","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# std multivariate independent normal latent space\nstd_img = tf.sigmoid( decoder.predict( [ std_sample, label_sample_vec ] ) ).numpy()\ndisplayImageBatch( std_img.reshape( -1, *IMAGE_SIZE[ :-1 ] ), label_sample )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# with total covariance matrix + mean in latent space\ntotal_sample = np.matmul( std_sample, totalCov ) + totalCenter\ntotal_img = tf.sigmoid( decoder.predict( [ std_sample, label_sample_vec ] ) ).numpy()\ndisplayImageBatch( total_img.reshape( -1, *IMAGE_SIZE[ :-1 ] ), label_sample )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# with per-class covariance matrix + mean in latent space\nclass_sample = np.zeros_like( std_sample )\nfor l in range( nClasses ):\n    pos = label_sample == l\n    class_sample[ pos ] = np.matmul( std_sample[ pos ], classCov[ l ] ) + classCenter[ l ]\nnp.matmul( std_sample, totalCov ) + totalCenter\nclass_img = tf.sigmoid( decoder.predict( [ std_sample, label_sample_vec ] ) ).numpy()\ndisplayImageBatch( class_img.reshape( -1, *IMAGE_SIZE[ :-1 ] ), label_sample )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Notebook by Mason Rumuly\n\nThanks to the Tensorflow tutorial on [convolutional variational auto encoders](https://www.tensorflow.org/tutorials/generative/cvae) on which the architecture of this model was based.\n\nIf you see something that could be improved, want further explaination, or just liked this notebook, please leave a comment!","execution_count":null}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}