exp_colored_mnist_P3.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"6\"\n",
    "import tensorflow as tf\n",
    "import logging\n",
    "tf.get_logger().setLevel(logging.ERROR)\n",
    "import tensorflow.keras.backend as K\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from tensorflow.keras import activations\n",
    "from tensorflow.keras.models import Model,Sequential\n",
    "from tensorflow.keras.layers import Dense,Conv2D,BatchNormalization,Lambda,Activation,Input,Flatten,Conv2DTranspose,concatenate,Reshape,LeakyReLU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set some options and read data based on selected ones"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'s_class_Isotropic_uniform_unsupervised_P3'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#constants\n",
    "IMG_DIM = 28\n",
    "# manual options\n",
    "sensitive = 'class' # 'color' or 'class' # class balanced left\n",
    "prior_dist = 0 # 0: simple gaussian 1: guassian mixture Note: if prior generator model enabled, this option is not used\n",
    "mnist_type = 'balanced' # 'balanced' 'unbalanced'\n",
    "supervised = False #False:Unsupervised (use input data as utility)\n",
    "data_path = \"./Colored_MNIST\"\n",
    "\n",
    "# automatically set variables based on selected options\n",
    "utility = 'class' if sensitive=='color' else 'color'\n",
    "prior_type = 'Isotropic' if prior_dist == 0 else 'Mixture'\n",
    "data_type = 'biased' if mnist_type=='unbalanced' else 'uniform'\n",
    "\n",
    "exp_info = f\"s_{sensitive}_{prior_type}_{data_type}_supervised\" if supervised else f\"s_{sensitive}_{prior_type}_{data_type}_unsupervised\"\n",
    "exp_info += \"_P3\"\n",
    "\n",
    "data = np.load(f\"{data_path}/colored_mnist_{mnist_type}.npz\")\n",
    "\n",
    "x_train = data['x_train']\n",
    "x_train = x_train.astype(np.float32) / 255.0\n",
    "s_train = data['s_train'] if(sensitive == 'color') else data['y_train']\n",
    "\n",
    "if not supervised:\n",
    "    u_train = x_train\n",
    "else:\n",
    "    u_train = data['s_train'] if(sensitive == 'class') else data['y_train']\n",
    "\n",
    "\n",
    "x_test = data['x_test']\n",
    "x_test = x_test.astype(np.float32) / 255.0\n",
    "s_test = data['s_test'] if(sensitive == 'color') else data['y_test']\n",
    "\n",
    "if not supervised:\n",
    "    u_test = x_test\n",
    "else:\n",
    "    u_test = data['s_test'] if(sensitive == 'class') else data['y_test']\n",
    "\n",
    "\n",
    "exp_info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# one hot S and U (if supervised) in order to ready them for the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def one_hot(labels, num_labels):\n",
    "    num_labels_data = labels.shape[0]\n",
    "    one_hot_encoding = np.zeros((num_labels_data,num_labels))\n",
    "    one_hot_encoding[np.arange(num_labels_data),labels] = 1\n",
    "    one_hot_encoding = np.reshape(one_hot_encoding, [-1, num_labels])\n",
    "    return one_hot_encoding\n",
    "\n",
    "DIM_S = 3 if(sensitive == 'color') else 10\n",
    "if supervised:\n",
    "    DIM_U = 10 if(sensitive == 'color') else 3\n",
    "else:\n",
    "    DIM_U = x_train[0].shape\n",
    "\n",
    "if supervised: u_train = one_hot(u_train, DIM_U).astype(np.float32)\n",
    "s_train = one_hot(s_train, DIM_S).astype(np.float32)\n",
    "if supervised: u_test = one_hot(u_test, DIM_U).astype(np.float32)\n",
    "s_test = one_hot(s_test, DIM_S).astype(np.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sampling(args):\n",
    "    \"\"\"Reparameterization trick by sampling from an isotropic unit Gaussian.\n",
    "\n",
    "    # Arguments\n",
    "        args (tensor): mean and log of variance of Q(z|X)\n",
    "\n",
    "    # Returns\n",
    "        z (tensor): sampled latent vector\n",
    "    \"\"\"\n",
    "\n",
    "    z_mean, z_log_var = args\n",
    "    batch = K.shape(z_mean)[0]\n",
    "    dim = K.int_shape(z_mean)[1]\n",
    "    # by default, random_normal has mean = 0 and std = 1.0\n",
    "    epsilon = K.random_normal(shape=(batch, dim))\n",
    "    return z_mean + K.exp(0.5 * z_log_var) * epsilon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_encoder(DIM_Z, input_x):\n",
    "    stride = 2\n",
    "#     input_x = Input( shape = [IMG_DIM,IMG_DIM,3], name=\"x\" )\n",
    "\n",
    "    #first hidden layer\n",
    "    x = Conv2D(64, 5, strides=stride, padding=\"same\", name=\"enc_h1\")(input_x)\n",
    "    x = BatchNormalization(name=\"enc_h1_normalized\")(x)\n",
    "    x = Activation(LeakyReLU(), name=\"enc_h1_activation\")(x)\n",
    "    #second hidden layer\n",
    "    x = Conv2D(128, 5, strides=stride, padding=\"same\", name=\"enc_h2\")(x)\n",
    "    x = BatchNormalization(name=\"enc_h2_normalized\")(x)\n",
    "    x = Activation(LeakyReLU(), name=\"enc_h2_activation\")(x)\n",
    "\n",
    "    shape = K.int_shape(x)\n",
    "    x = Flatten()(x)\n",
    "    x = Dense(DIM_Z*4, name=\"enc_dense_1\")(x)\n",
    "    x = BatchNormalization()(x)\n",
    "    x = Activation('tanh')(x)\n",
    "\n",
    "    z_mean = Dense(DIM_Z, name=\"z_mean\")(x)\n",
    "    z_log_sigma_sq = Dense(DIM_Z, name=\"z_sigma\")(x)\n",
    "    z = Lambda(sampling, output_shape=DIM_Z, name='z')([z_mean, z_log_sigma_sq])\n",
    "    \n",
    "#     prior_loss = K.mean(-0.5 * K.sum(1 + z_log_sigma_sq - K.square(z_mean) - K.exp(z_log_sigma_sq), axis=-1))\n",
    "    \n",
    "    encoder = Model(input_x, z, name = \"Encoder\")\n",
    "#     encoder.add_loss((alpha+beta) * prior_loss)\n",
    "    \n",
    "    return (encoder, z_mean, z_log_sigma_sq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_uncertainty_model(DIM_Z):\n",
    "    model = Sequential(name=\"Uncertainty_Decoder\")\n",
    "    model.add(Dense(DIM_Z * 4, input_dim=DIM_Z))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    model.add(Dense(DIM_S, activation='softmax', name=\"s_hat\"))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "if supervised:\n",
    "    def get_utility_model(DIM_Z):\n",
    "        model = Sequential(name=\"Utility_Decoder_supervised\")\n",
    "        model.add(Dense(DIM_Z * 4, input_dim=DIM_Z))\n",
    "        model.add(BatchNormalization())\n",
    "        model.add(LeakyReLU(alpha=0.2))\n",
    "        model.add(Dense(DIM_U, activation='softmax', name=\"u_hat\"))\n",
    "        return model\n",
    "else:\n",
    "    def get_utility_model(DIM_Z):\n",
    "        stride = 2\n",
    "        model = Sequential(name=\"Utility_Decoder_unsupervised\")\n",
    "        model.add(Dense(7*7*128, input_dim=DIM_Z))\n",
    "        model.add(BatchNormalization())\n",
    "        model.add(LeakyReLU(alpha=0.2))\n",
    "        \n",
    "        model.add(Reshape((7,7,128)))\n",
    "        \n",
    "        model.add(Conv2DTranspose(64, 5, strides=stride, padding=\"same\"))\n",
    "        model.add(BatchNormalization())\n",
    "        model.add(LeakyReLU(alpha=0.2))\n",
    "        \n",
    "        model.add(Conv2DTranspose(3, 5, strides=stride, padding=\"same\"))\n",
    "        model.add(Activation(\"sigmoid\"))\n",
    "        return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_z_discriminator(DIM_Z):\n",
    "    model = Sequential(name=\"Latent_Space_Discriminator\")\n",
    "    \n",
    "    model.add(Dense(512, input_dim=DIM_Z))\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    \n",
    "    model.add(Dense(128))\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    \n",
    "    model.add(Dense(1, activation=\"sigmoid\"))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_utility_discriminator():\n",
    "    model = Sequential(name=\"Attribute_Class_Discriminator\")\n",
    "\n",
    "    model.add(Dense(DIM_U * 8, input_dim=DIM_U))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    \n",
    "    model.add(Dense(DIM_U * 8))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    \n",
    "    model.add(Dense(1, activation=\"sigmoid\"))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_visible_space_discriminator():\n",
    "    stride = 2\n",
    "    model = Sequential(name=\"Visible_Space_Discriminator\")\n",
    "    \n",
    "    model.add(Input(shape = [IMG_DIM,IMG_DIM,3]))\n",
    "    \n",
    "    model.add(Conv2D(32, 3, strides=stride, padding=\"same\"))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    \n",
    "    model.add(Conv2D(64, 3, strides=stride, padding=\"same\"))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    \n",
    "    model.add(Conv2D(128, 3, strides=stride, padding=\"same\"))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    \n",
    "    model.add(Flatten())\n",
    "    \n",
    "    model.add(Dense(128))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    \n",
    "    model.add(Dense(1, activation=\"sigmoid\"))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_prior_generator(DIM_Z, noise_dim=50):\n",
    "    model = Sequential(name=\"Prior_Distribution_Generator\")\n",
    "    \n",
    "    model.add(Dense(noise_dim*2, input_dim=noise_dim))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(Activation(\"relu\"))\n",
    "    \n",
    "    model.add(Dense(noise_dim))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(Activation(\"relu\"))\n",
    "    \n",
    "    model.add(Dense(DIM_Z))\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_uncertainty_discriminator():\n",
    "    model = Sequential(name=\"Uncertainty_Discriminator\")\n",
    "\n",
    "    model.add(Dense(DIM_S * 8, input_dim=DIM_S))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    \n",
    "    model.add(Dense(DIM_S * 8))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    \n",
    "    model.add(Dense(1, activation=\"sigmoid\"))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Putting all togther - define CLUB model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "z_gen_enabled = True #if False, the algorithm sample prior from Normal Distribution\n",
    "## utility and reconstruction losses\n",
    "if supervised:\n",
    "    def loss_u(gamma):\n",
    "        def loss(u_true, u_pred):\n",
    "            return gamma * K.mean(K.sum(K.square(u_true-u_pred), axis=-1))\n",
    "        return loss\n",
    "else:\n",
    "    def loss_u(gamma):\n",
    "        def loss(u_true, u_pred):\n",
    "            return gamma * K.mean(K.sum(K.square(u_true-u_pred), axis=(1,2,3)))\n",
    "        return loss\n",
    "\n",
    "def loss_s(lam):\n",
    "    def loss(s_true, s_pred):\n",
    "        return lam * -K.mean(K.sum(K.square(s_true-s_pred), axis=-1))\n",
    "    return loss\n",
    "\n",
    "bce = tf.keras.losses.BinaryCrossentropy()\n",
    "# Weighted cross-entropy loss\n",
    "def loss_wce(coef):\n",
    "    def loss(y, y_pred):\n",
    "         return coef * bce(y, y_pred)\n",
    "    return loss\n",
    "\n",
    "def get_full_model(DIM_Z, lam, gamma, learning_rate=0.0001, dim_noise = 100):\n",
    "    ########## Inputs\n",
    "    input_x = Input( shape=[IMG_DIM,IMG_DIM,3], name=\"x\" )\n",
    "    input_z = Input( shape = (DIM_Z,), name=\"z\" )\n",
    "    input_noise = Input( shape = (dim_noise,), name=\"Noise\")\n",
    "\n",
    "    ########## Define AE: Encoder, Utility Decoder and Uncertainty Decoder\n",
    "    encoder,z_mean,z_log_sigma_sq = get_encoder(DIM_Z, input_x)\n",
    "    uncertainty_decoder = get_uncertainty_model(DIM_Z)\n",
    "    utility_decoder = get_utility_model(DIM_Z)\n",
    "\n",
    "    z = encoder(input_x)\n",
    "    s_hat = uncertainty_decoder(z)\n",
    "    u_hat = utility_decoder(z)\n",
    "\n",
    "    autoencoder = Model(input_x, [s_hat, u_hat], name=\"CLUB_Autoencoder\")\n",
    "    prior_loss = K.mean(-0.5 * K.sum(1 + z_log_sigma_sq - K.square(z_mean) - K.exp(z_log_sigma_sq), axis=-1))\n",
    "    prior_loss = tf.identity(prior_loss, name=\"kl_loss\")\n",
    "    autoencoder.add_loss(prior_loss)\n",
    "    autoencoder.compile(loss=[loss_s(lam), loss_u(gamma)], optimizer=tf.keras.optimizers.Adam(lr=learning_rate*5))\n",
    "\n",
    "    ########## Define Latent Space Discriminator\n",
    "    z_discriminator = get_z_discriminator(DIM_Z)\n",
    "    z_discriminator.compile(loss=loss_wce(0.1), optimizer=tf.keras.optimizers.Adam(lr=learning_rate))\n",
    "\n",
    "    ##########\n",
    "    #z_discriminator should train separately\n",
    "    z_discriminator.trainable = False\n",
    "    prior_generator = get_prior_generator(DIM_Z, dim_noise)\n",
    "    prior_gen_zdiscriminator = Model(input_noise, z_discriminator(prior_generator(input_noise)), name=\"CLUB_generator_zdiscriminator\")\n",
    "    prior_gen_zdiscriminator.compile(loss=loss_wce(-0.1), optimizer=tf.keras.optimizers.Adam(lr=learning_rate))\n",
    "\n",
    "    ##########\n",
    "    encoder_zdiscriminator = Model(input_x, z_discriminator(encoder(input_x)), name=\"CLUB_encoder_zdiscriminator\")\n",
    "    encoder_zdiscriminator.compile(loss=loss_wce(-0.1), optimizer=tf.keras.optimizers.Adam(lr=learning_rate))\n",
    "\n",
    "    ##########\n",
    "    u_dircriminator = get_utility_discriminator() if supervised else get_visible_space_discriminator()\n",
    "    u_dircriminator.compile(loss=loss_wce(0.1*gamma), optimizer=tf.keras.optimizers.Adam(lr=learning_rate))\n",
    "\n",
    "    ##########\n",
    "    #u_dircriminator should train separately\n",
    "    u_dircriminator.trainable = False\n",
    "\n",
    "    if z_gen_enabled:    \n",
    "        decoder_udiscriminator = Model(input_noise, u_dircriminator(utility_decoder(prior_generator(input_noise))), name=\"CLUB_decoder_udiscriminator\")\n",
    "    else:\n",
    "        decoder_udiscriminator = Model(input_z, u_dircriminator(utility_decoder(input_z)), name=\"CLUB_decoder_udiscriminator\")\n",
    "\n",
    "    decoder_udiscriminator.compile(loss=loss_wce(-0.1*gamma), optimizer=tf.keras.optimizers.Adam(lr=learning_rate))\n",
    "\n",
    "    ##########\n",
    "    s_dircriminator = get_uncertainty_discriminator()\n",
    "    s_dircriminator.compile(loss=loss_wce(-0.1*lam), optimizer=tf.keras.optimizers.Adam(lr=learning_rate))\n",
    "\n",
    "    ##########\n",
    "    #s_dircriminator should train separately\n",
    "    s_dircriminator.trainable = False\n",
    "    if z_gen_enabled:\n",
    "        decoder_sdiscriminator = Model(input_noise, s_dircriminator(uncertainty_decoder(prior_generator(input_noise))), name=\"CLUB_decoder_sdiscriminator\")\n",
    "    else:\n",
    "        decoder_sdiscriminator = Model(input_z, s_dircriminator(uncertainty_decoder(input_z)), name=\"CLUB_decoder_sdiscriminator\")\n",
    "\n",
    "    decoder_sdiscriminator.compile(loss=loss_wce(0.1 * lam), optimizer=tf.keras.optimizers.Adam(lr=learning_rate))\n",
    "                                   \n",
    "    return encoder, uncertainty_decoder, utility_decoder, autoencoder, z_discriminator, prior_generator, prior_gen_zdiscriminator, encoder_zdiscriminator, u_dircriminator, decoder_udiscriminator, s_dircriminator, decoder_sdiscriminator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for test\n",
    "# encoder, uncertainty_decoder, utility_decoder, autoencoder, z_discriminator, prior_generator, prior_gen_zdiscriminator, encoder_zdiscriminator, u_dircriminator, decoder_udiscriminator, s_dircriminator, decoder_sdiscriminator= get_full_model(32,0.1,0.1)\n",
    "# # tf.keras.utils.plot_model(decoder_udiscriminator, show_shapes=True)\n",
    "# tf.keras.utils.plot_model(autoencoder,show_shapes=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pre-training implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pre_train_ae(dim_z, lam, gamma, force_train=False, max_itr=50, batch_size=512, verbose=2):\n",
    "    info_str = f\"d_{dim_z}_lambda_{lam}_gamma_{gamma}_{exp_info}\"\n",
    "    if not os.path.exists(f\"saved_models/mnist_pretrain_{info_str}.h5\") or force_train:\n",
    "        print(f\"Pre-Training with {info_str}\")\n",
    "        autoencoder.fit(x_train, [s_train, u_train], batch_size=batch_size, epochs=max_itr, shuffle=True, verbose=verbose)\n",
    "        autoencoder.save_weights(f\"saved_models/mnist_pretrain_{info_str}.h5\")\n",
    "    else:\n",
    "        print(f\"Loading model from file with {info_str}\")\n",
    "        autoencoder.load_weights(f\"saved_models/mnist_pretrain_{info_str}.h5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Full model training loop "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(sum([np.sum(w) for w in pib_autoencoder.get_weights()]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.sampler import gaussian, gaussian_mixture\n",
    "\n",
    "def sample_prior(latent_dim, batch_size):\n",
    "    if prior_dist == 0:  \n",
    "        return gaussian(batch_size, latent_dim)\n",
    "    elif prior_dist == 1:\n",
    "        return gaussian_mixture(batch_size, latent_dim, num_labels=min(10, DIM_U))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main_train(dim_z, max_itr=2000, batch_size=1024, z_disc_enabled=True,u_disc_enabled=True,s_disc_enabled=True,verbose=0,dim_noise=100):\n",
    "    ones = np.ones((batch_size, 1))\n",
    "    zeros = np.zeros((batch_size, 1))\n",
    "    start_time = time.time()\n",
    "    for epoch in range(max_itr):\n",
    "        start_time_epoch = time.time()\n",
    "        # Select a random batch of images\n",
    "        idx = np.random.randint(0, x_train.shape[0], batch_size)\n",
    "        x = x_train[idx]\n",
    "        s = s_train[idx]\n",
    "        u = u_train[idx]\n",
    "        # ---------------------\n",
    "        #  1- Train the Encoder, Utility Decoder, Uncertainty Decoder\n",
    "        # ---------------------\n",
    "        ae_loss = autoencoder.train_on_batch(x, [s, u])\n",
    "\n",
    "        if z_disc_enabled:\n",
    "        # ---------------------\n",
    "        #  2- Train the Latent Space Discriminator\n",
    "        # ---------------------\n",
    "\n",
    "            if z_gen_enabled:\n",
    "                noise = np.random.uniform(-1.0, 1.0, size=[batch_size, dim_noise])\n",
    "                latent_prior = prior_generator(noise)\n",
    "            else:\n",
    "                latent_prior = sample_prior(dim_z, batch_size)\n",
    "\n",
    "            idx = np.random.randint(0, x_train.shape[0], batch_size)\n",
    "            x = x_train[idx]\n",
    "            latent_enc = encoder(x)\n",
    "\n",
    "            d_loss_prior = z_discriminator.train_on_batch(latent_prior, zeros)\n",
    "            d_loss_enc = z_discriminator.train_on_batch(latent_enc, ones)\n",
    "            dz_loss = d_loss_prior + d_loss_enc\n",
    "\n",
    "            # ---------------------\n",
    "            # 3- Train the Encoder and Prior Distribution Generator Adversarially\n",
    "            # ---------------------\n",
    "            prior_loss = 0.\n",
    "            if z_gen_enabled:\n",
    "                noise = np.random.uniform(-1.0, 1.0, size=[batch_size, dim_noise])\n",
    "                prior_loss = prior_gen_zdiscriminator.train_on_batch(noise, zeros)\n",
    "\n",
    "            idx = np.random.randint(0, x_train.shape[0], batch_size)\n",
    "            x = x_train[idx]\n",
    "            edz_loss = encoder_zdiscriminator.train_on_batch(x, ones)\n",
    "        else:\n",
    "            dz_loss = 0.\n",
    "            prior_loss = 0.\n",
    "            edz_loss = 0.\n",
    "\n",
    "        if u_disc_enabled:\n",
    "        # ---------------------\n",
    "        #  4- Train Visible_Space/Attribute_Class Discriminator \n",
    "        # ---------------------\n",
    "\n",
    "            if z_gen_enabled:\n",
    "                noise = np.random.uniform(-1.0, 1.0, size=[batch_size, dim_noise])\n",
    "                latent = prior_generator(noise)\n",
    "            else:\n",
    "                latent = sample_prior(dim_z, batch_size)\n",
    "            u_dec = utility_decoder(latent)\n",
    "\n",
    "            idx = np.random.randint(0, x_train.shape[0], batch_size)\n",
    "            u = u_train[idx]\n",
    "\n",
    "            d_loss_real = u_dircriminator.train_on_batch(u, ones)\n",
    "            d_loss_fake = u_dircriminator.train_on_batch(u_dec, zeros)\n",
    "            du_loss = d_loss_real + d_loss_fake\n",
    "\n",
    "        # ---------------------\n",
    "        #  5- Train the Prior Distribution Generator and Utility Decoder Adversarially\n",
    "        # ---------------------\n",
    "            if z_gen_enabled:\n",
    "                noise = np.random.uniform(-1.0, 1.0, size=[batch_size, dim_noise])\n",
    "                gdu_loss = decoder_udiscriminator.train_on_batch(noise, zeros)\n",
    "            else:    \n",
    "                latent = sample_prior(dim_z, batch_size)\n",
    "                gdu_loss = decoder_udiscriminator.train_on_batch(latent, zeros)\n",
    "        else:\n",
    "            gdu_loss = 0.\n",
    "            du_loss = 0.\n",
    "\n",
    "        # ---------------------\n",
    "        #  6- Train the Uncertainty Discriminator\n",
    "        # ---------------------\n",
    "        if s_disc_enabled:\n",
    "            \n",
    "            if z_gen_enabled:\n",
    "                noise = np.random.uniform(-1.0, 1.0, size=[batch_size, dim_noise])\n",
    "                latent = prior_generator(noise)\n",
    "            else:\n",
    "                latent = sample_prior(dim_z, batch_size)\n",
    "            s_dec = uncertainty_decoder(latent)\n",
    "\n",
    "            idx = np.random.randint(0, x_train.shape[0], batch_size)\n",
    "            s = s_train[idx]\n",
    "\n",
    "            d_loss_1 = s_dircriminator.train_on_batch(s, ones)\n",
    "            d_loss_2 = s_dircriminator.train_on_batch(s_dec, zeros)\n",
    "            ds_loss = d_loss_1 + d_loss_2\n",
    "        # ---------------------\n",
    "        #  7- Train the Prior Distribution Generator and Uncertainty Decoder\n",
    "        # ---------------------\n",
    "            if z_gen_enabled:\n",
    "                noise = np.random.uniform(-1.0, 1.0, size=[batch_size, dim_noise])\n",
    "                gds_loss = decoder_sdiscriminator.train_on_batch(noise, zeros)\n",
    "            else:    \n",
    "                latent = sample_prior(dim_z, batch_size)\n",
    "                gds_loss = decoder_sdiscriminator.train_on_batch(latent, zeros)\n",
    "        else:\n",
    "            ds_loss = 0.\n",
    "            gds_loss = 0.\n",
    "        \n",
    "        # ---------------------\n",
    "        #  Print stats info\n",
    "        # ---------------------\n",
    "        if verbose != 0 and epoch % 50 == 0:\n",
    "            print(f\"{epoch}, s:{ae_loss[1]:.4f}, u:{ae_loss[2]:.4f}, dz:{dz_loss:.4f}, edz:{edz_loss:.4f}, prior:{prior_loss:.4f}, du:{du_loss:.4f}, gdu:{gdu_loss:.4f}, ds:{ds_loss:.4f}, gds:{gds_loss:.4f}\")\n",
    "            print(f\"One epoch execution time: {(time.time() - start_time_epoch):.5} seconds\")\n",
    "\n",
    "    total_time = (time.time() - start_time) / 60\n",
    "    print(f\"Total Execution Time: {total_time:.4f} minutes\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Experiment 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pre-Training with d_8_lambda_0.0001_gamma_0.0001_s_class_Isotropic_uniform_unsupervised_P3\n"
     ]
    },
    {
     "ename": "UnknownError",
     "evalue": " Failed to get convolution algorithm. This is probably because cuDNN failed to initialize, so try looking to see if a warning log message was printed above.\n\t [[node CLUB_Autoencoder/enc_h1/Conv2D (defined at <ipython-input-15-7b6f51b1c950>:5) ]] [Op:__inference_train_function_5157]\n\nFunction call stack:\ntrain_function\n",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mUnknownError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-19-643fede841d2>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     20\u001b[0m             \u001b[0;31m#Pre training\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m             \u001b[0mencoder\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0muncertainty_decoder\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mutility_decoder\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mautoencoder\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mz_discriminator\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprior_generator\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprior_gen_zdiscriminator\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mencoder_zdiscriminator\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mu_dircriminator\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdecoder_udiscriminator\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0ms_dircriminator\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdecoder_sdiscriminator\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_full_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdimz\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlam\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgamma\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 22\u001b[0;31m             \u001b[0mpre_train_ae\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdimz\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlam\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgamma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mforce_train\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmax_itr\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m50\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1024\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mverbose\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     23\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     24\u001b[0m             \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Full model training\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-15-7b6f51b1c950>\u001b[0m in \u001b[0;36mpre_train_ae\u001b[0;34m(dim_z, lam, gamma, force_train, max_itr, batch_size, verbose)\u001b[0m\n\u001b[1;32m      3\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexists\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"saved_models/mnist_pretrain_{info_str}.h5\"\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mforce_train\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m         \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"Pre-Training with {info_str}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m         \u001b[0mautoencoder\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx_train\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0ms_train\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mu_train\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mbatch_size\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mepochs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmax_itr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mshuffle\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mverbose\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mverbose\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      6\u001b[0m         \u001b[0mautoencoder\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave_weights\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"saved_models/mnist_pretrain_{info_str}.h5\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/tensorflow-24/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py\u001b[0m in \u001b[0;36mfit\u001b[0;34m(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, max_queue_size, workers, use_multiprocessing)\u001b[0m\n\u001b[1;32m   1098\u001b[0m                 _r=1):\n\u001b[1;32m   1099\u001b[0m               \u001b[0mcallbacks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_train_batch_begin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1100\u001b[0;31m               \u001b[0mtmp_logs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0miterator\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1101\u001b[0m               \u001b[0;32mif\u001b[0m \u001b[0mdata_handler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshould_sync\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1102\u001b[0m                 \u001b[0mcontext\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0masync_wait\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/tensorflow-24/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwds)\u001b[0m\n\u001b[1;32m    826\u001b[0m     \u001b[0mtracing_count\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexperimental_get_tracing_count\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    827\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mtrace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTrace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_name\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mtm\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 828\u001b[0;31m       \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    829\u001b[0m       \u001b[0mcompiler\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"xla\"\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_experimental_compile\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;34m\"nonXla\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    830\u001b[0m       \u001b[0mnew_tracing_count\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexperimental_get_tracing_count\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/tensorflow-24/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py\u001b[0m in \u001b[0;36m_call\u001b[0;34m(self, *args, **kwds)\u001b[0m\n\u001b[1;32m    886\u001b[0m         \u001b[0;31m# Lifting succeeded, so variables are initialized and we can run the\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    887\u001b[0m         \u001b[0;31m# stateless function.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 888\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_stateless_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    889\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    890\u001b[0m       \u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfiltered_flat_args\u001b[0m \u001b[0;34m=\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m\\\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/tensorflow-24/lib/python3.7/site-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   2941\u001b[0m        filtered_flat_args) = self._maybe_define_function(args, kwargs)\n\u001b[1;32m   2942\u001b[0m     return graph_function._call_flat(\n\u001b[0;32m-> 2943\u001b[0;31m         filtered_flat_args, captured_inputs=graph_function.captured_inputs)  # pylint: disable=protected-access\n\u001b[0m\u001b[1;32m   2944\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2945\u001b[0m   \u001b[0;34m@\u001b[0m\u001b[0mproperty\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/tensorflow-24/lib/python3.7/site-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36m_call_flat\u001b[0;34m(self, args, captured_inputs, cancellation_manager)\u001b[0m\n\u001b[1;32m   1917\u001b[0m       \u001b[0;31m# No tape is watching; skip to running the function.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1918\u001b[0m       return self._build_call_outputs(self._inference_function.call(\n\u001b[0;32m-> 1919\u001b[0;31m           ctx, args, cancellation_manager=cancellation_manager))\n\u001b[0m\u001b[1;32m   1920\u001b[0m     forward_backward = self._select_forward_and_backward_functions(\n\u001b[1;32m   1921\u001b[0m         \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/tensorflow-24/lib/python3.7/site-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36mcall\u001b[0;34m(self, ctx, args, cancellation_manager)\u001b[0m\n\u001b[1;32m    558\u001b[0m               \u001b[0minputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    559\u001b[0m               \u001b[0mattrs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mattrs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 560\u001b[0;31m               ctx=ctx)\n\u001b[0m\u001b[1;32m    561\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    562\u001b[0m           outputs = execute.execute_with_cancellation(\n",
      "\u001b[0;32m~/anaconda3/envs/tensorflow-24/lib/python3.7/site-packages/tensorflow/python/eager/execute.py\u001b[0m in \u001b[0;36mquick_execute\u001b[0;34m(op_name, num_outputs, inputs, attrs, ctx, name)\u001b[0m\n\u001b[1;32m     58\u001b[0m     \u001b[0mctx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mensure_initialized\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     59\u001b[0m     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,\n\u001b[0;32m---> 60\u001b[0;31m                                         inputs, attrs, num_outputs)\n\u001b[0m\u001b[1;32m     61\u001b[0m   \u001b[0;32mexcept\u001b[0m \u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_NotOkStatusException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     62\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mname\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mUnknownError\u001b[0m:  Failed to get convolution algorithm. This is probably because cuDNN failed to initialize, so try looking to see if a warning log message was printed above.\n\t [[node CLUB_Autoencoder/enc_h1/Conv2D (defined at <ipython-input-15-7b6f51b1c950>:5) ]] [Op:__inference_train_function_5157]\n\nFunction call stack:\ntrain_function\n"
     ]
    }
   ],
   "source": [
    "from utils.mine import MINE\n",
    "lam_list = [0.0001, 0.1, 10, 100, 1000]\n",
    "gamma_list = [0.0001, 0.1, 10, 100, 1000]\n",
    "DIM_Z = [8, 64]\n",
    "\n",
    "util_acc_tr = np.zeros((len(DIM_Z),len(lam_list),len(gamma_list)))\n",
    "util_acc_ts = np.zeros((len(DIM_Z),len(lam_list),len(gamma_list)))\n",
    "sens_acc_tr = np.zeros((len(DIM_Z),len(lam_list),len(gamma_list)))\n",
    "sens_acc_ts = np.zeros((len(DIM_Z),len(lam_list),len(gamma_list)))\n",
    "\n",
    "mi_u_ts = np.zeros((len(DIM_Z),len(lam_list),len(gamma_list)))\n",
    "mi_u_tr = np.zeros((len(DIM_Z),len(lam_list),len(gamma_list)))\n",
    "mi_s_ts = np.zeros((len(DIM_Z),len(lam_list),len(gamma_list)))\n",
    "mi_s_tr = np.zeros((len(DIM_Z),len(lam_list),len(gamma_list)))\n",
    "\n",
    "\n",
    "for i, dimz in enumerate(DIM_Z):\n",
    "    for j, lam in enumerate(lam_list):\n",
    "        for k, gamma in enumerate(gamma_list):\n",
    "            #Pre training\n",
    "            encoder, uncertainty_decoder, utility_decoder, autoencoder, z_discriminator, prior_generator, prior_gen_zdiscriminator, encoder_zdiscriminator, u_dircriminator, decoder_udiscriminator, s_dircriminator, decoder_sdiscriminator = get_full_model(dimz, lam, gamma)\n",
    "            pre_train_ae(dimz, lam, gamma, force_train=False, max_itr=50, batch_size=1024, verbose=0)\n",
    "            \n",
    "            print(\"Full model training\")\n",
    "            main_train(dimz, max_itr=100, batch_size=1024,verbose=1)\n",
    "\n",
    "            print(\"Evaluate performance\")\n",
    "            z_test = encoder.predict(x_test)\n",
    "            z_train = encoder.predict(x_train)\n",
    "            \n",
    "            u_test_hat = utility_decoder.predict(z_test)\n",
    "            u_train_hat = utility_decoder.predict(z_train)\n",
    "            \n",
    "            if supervised:\n",
    "                util_acc_ts[i][j][k] = np.mean(np.argmax(u_test_hat,axis=1)==np.argmax(u_test,axis=1)) * 100\n",
    "                util_acc_tr[i][j][k] = np.mean(np.argmax(u_train_hat,axis=1)==np.argmax(u_train,axis=1)) * 100\n",
    "            else:\n",
    "                util_acc_ts[i][j][k] = np.mean(np.sum(np.square(u_test_hat-u_test), axis=(1,2,3)))\n",
    "                util_acc_tr[i][j][k] = np.mean(np.sum(np.square(u_train_hat-u_train), axis=(1,2,3)))\n",
    "            \n",
    "            s_test_hat = uncertainty_decoder.predict(z_test)\n",
    "            sens_acc_ts[i][j][k] = np.mean(np.argmax(s_test_hat,axis=1)==np.argmax(s_test,axis=1)) * 100\n",
    "            \n",
    "            s_train_hat = uncertainty_decoder.predict(z_train)\n",
    "            sens_acc_tr[i][j][k] = np.mean(np.argmax(s_train_hat,axis=1)==np.argmax(s_train,axis=1)) * 100\n",
    "            \n",
    "            if supervised:\n",
    "                print(\"Evaluate Mutual Information I(Z;U)\")            \n",
    "                mine = MINE(x_dim=dimz, y_dim=DIM_U)\n",
    "                _, mi_u_ts[i][j][k] = mine.fit(z_test, u_test, epochs=200, batch_size=1024, verbose=0)\n",
    "                _, mi_u_tr[i][j][k] = mine.fit(z_train, u_train, epochs=250, batch_size=1024, verbose=0)\n",
    "            \n",
    "            print(\"Evaluate Mutual Information I(Z;S)\")\n",
    "            mine = MINE(x_dim=dimz, y_dim=DIM_S)\n",
    "            _, mi_s_ts[i][j][k] = mine.fit(z_test, s_test, epochs=250, batch_size=1024, verbose=0)\n",
    "            _, mi_s_tr[i][j][k] = mine.fit(z_train, s_train, epochs=250, batch_size=1024, verbose=0) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot results\n",
    "## Plot accuracy (for each z we save one fig for all lambda on the Test data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams.update({'font.size': 16})\n",
    "plt.rcParams['figure.dpi'] = 300\n",
    "linestyles = ['-', '--', '-.', ':']\n",
    "for i in range(util_acc_ts.shape[0]):\n",
    "    plt.figure(figsize=(15,8))\n",
    "#     if supervised:\n",
    "#         plt.title(r'Utility and Sensitive attributes accuracy with $\\gamma$ $\\in$ ' + f'{gamma_list} {exp_info}_d={DIM_Z[i]}', fontsize=16)\n",
    "#     else:\n",
    "#         plt.title(r'Utility attribute Mean Square Error with $\\gamma$ $\\in$ ' + f'{gamma_list} {exp_info}_d={DIM_Z[i]}', fontsize=16)\n",
    "        \n",
    "    for j in range(util_acc_ts.shape[1]):\n",
    "        if supervised:\n",
    "            plt.plot(util_acc_ts[i][j], label=r'Utility Att Acc., Test, $\\lambda=$' + f\"{lam_list[j]}\", linestyle=linestyles[0], linewidth=2)\n",
    "            plt.plot(sens_acc_ts[i][j], label=r'Sensitive Att Acc., Test, $\\lambda=$' + f\"{lam_list[j]}\", linestyle=linestyles[0], linewidth=2)\n",
    "        else:\n",
    "            plt.plot(util_acc_ts[i][j], label=r'Utility Att MSE., Test, $\\lambda=$' + f\"{lam_list[j]}\", linestyle=linestyles[0], linewidth=2)\n",
    "\n",
    "\n",
    "    plt.legend(prop={'size': 16})\n",
    "    plt.grid()\n",
    "    plt.xticks(list(range(len(gamma_list))), gamma_list, fontsize=18, rotation=90)\n",
    "    plt.xlabel(r'$\\gamma$', fontsize=24)\n",
    "    if supervised:\n",
    "        plt.ylabel(r'Accuracy on $\\mathbf{U}$ and $\\mathbf{S}$', fontsize=24)\n",
    "        plt.savefig(f'./saved_figures/chart_colored_mnist_acc_d_{DIM_Z[i]}_{exp_info}.svg', format='svg', bbox_inches='tight')\n",
    "        plt.savefig(f'./saved_figures/chart_colored_mnist_acc_d_{DIM_Z[i]}_{exp_info}.png', format='png', bbox_inches='tight')\n",
    "    else:\n",
    "        plt.ylabel(r'MSE on $\\mathbf{U}$', fontsize=24)\n",
    "        plt.savefig(f'./saved_figures/chart_colored_mnist_acc_mse_d_{DIM_Z[i]}_{exp_info}.svg', format='svg', bbox_inches='tight')\n",
    "        plt.savefig(f'./saved_figures/chart_colored_mnist_acc_mse_d_{DIM_Z[i]}_{exp_info}.png', format='png', bbox_inches='tight')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot accuracy (for each z and lambda we save and plot a figure on the Test and Train data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams.update({'font.size': 16})\n",
    "plt.rcParams['figure.dpi'] = 300\n",
    "linestyles = ['-', '--', '-.', ':']\n",
    "for i in range(util_acc_ts.shape[0]):\n",
    "    for j in range(util_acc_ts.shape[1]):\n",
    "        plt.figure(figsize=(15,8))\n",
    "        if supervised:\n",
    "#             plt.title(r'Utility and Sensitive attributes accuracy with $\\gamma$ $\\in$ ' + f'{gamma_list} {exp_info}_d={DIM_Z[i]}_lam={lam_list[j]}', fontsize=16)\n",
    "            plt.plot(util_acc_ts[i][j], label=r'Utility Att Acc., Test $\\lambda=$' + f\"{lam_list[j]}\", linestyle=linestyles[0], linewidth=2)\n",
    "            plt.plot(util_acc_tr[i][j], label=r'Utility Att Acc., Train $\\lambda=$' + f\"{lam_list[j]}\", linestyle=linestyles[1], linewidth=3)\n",
    "\n",
    "            plt.plot(sens_acc_ts[i][j], label=r'Sensitive Att Acc., Test $\\lambda=$' + f\"{lam_list[j]}\", linestyle=linestyles[0], linewidth=2)\n",
    "            plt.plot(sens_acc_tr[i][j], label=r'Sensitive Att Acc., Train $\\lambda=$' + f\"{lam_list[j]}\", linestyle=linestyles[1], linewidth=3)\n",
    "        else:\n",
    "#             plt.title(r'Utility attribute Mean Square Error with $\\gamma$ $\\in$ ' + f'{gamma_list} {exp_info}_d={DIM_Z[i]}_lam={lam_list[j]}', fontsize=16)\n",
    "            plt.plot(util_acc_ts[i][j], label=r'Utility Att MSE., Test', linestyle=linestyles[0], linewidth=2)\n",
    "            plt.plot(util_acc_tr[i][j], label=r'Utility Att MSE., Train', linestyle=linestyles[1], linewidth=3)\n",
    "\n",
    "\n",
    "        plt.legend(prop={'size': 16})\n",
    "        plt.grid()\n",
    "        plt.xticks(list(range(len(gamma_list))), gamma_list, fontsize=18, rotation=90)\n",
    "        plt.xlabel(r'$\\gamma$', fontsize=24)\n",
    "        if supervised:\n",
    "            plt.ylabel(r'Accuracy on $\\mathbf{U}$ and $\\mathbf{S}$', fontsize=24)\n",
    "            plt.savefig(f'./saved_figures/chart_colored_mnist_acc_d_{DIM_Z[i]}_lam_{lam_list[j]}_{exp_info}.svg', format='svg', bbox_inches='tight')\n",
    "            plt.savefig(f'./saved_figures/chart_colored_mnist_acc_d_{DIM_Z[i]}_lam_{lam_list[j]}_{exp_info}.png', format='png', bbox_inches='tight')\n",
    "        else:\n",
    "            plt.ylabel(r'MSE on $\\mathbf{U}$', fontsize=24)\n",
    "            plt.savefig(f'./saved_figures/chart_colored_mnist_acc_mse_d_{DIM_Z[i]}_lam_{lam_list[j]}_{exp_info}.svg', format='svg', bbox_inches='tight')\n",
    "            plt.savefig(f'./saved_figures/chart_colored_mnist_acc_mse_d_{DIM_Z[i]}_lam_{lam_list[j]}_{exp_info}.png', format='png', bbox_inches='tight')\n",
    "        plt.show()\n",
    "\n",
    "if not supervised:\n",
    "    for i in range(util_acc_ts.shape[0]):\n",
    "        for j in range(util_acc_ts.shape[1]):\n",
    "            plt.figure(figsize=(15,8))\n",
    "            plt.plot(sens_acc_ts[i][j], label=r'Sensitive Att Acc., Test $\\lambda=$' + f\"{lam_list[j]}\", linestyle=linestyles[0], linewidth=2)\n",
    "            plt.plot(sens_acc_tr[i][j], label=r'Sensitive Att Acc., Train $\\lambda=$' + f\"{lam_list[j]}\", linestyle=linestyles[1], linewidth=3)\n",
    "\n",
    "\n",
    "            plt.legend(prop={'size': 16})\n",
    "            plt.grid()\n",
    "            plt.xticks(list(range(len(gamma_list))), gamma_list, fontsize=18, rotation=90)\n",
    "            plt.xlabel(r'$\\gamma$', fontsize=24)\n",
    "            plt.ylabel(r'Accuracy on $\\mathbf{S}$', fontsize=24)\n",
    "            plt.savefig(f'./saved_figures/chart_colored_mnist_acc_d_{DIM_Z[i]}_lam_{lam_list[j]}_{exp_info}.svg', format='svg', bbox_inches='tight')\n",
    "            plt.savefig(f'./saved_figures/chart_colored_mnist_acc_d_{DIM_Z[i]}_lam_{lam_list[j]}_{exp_info}.png', format='png', bbox_inches='tight')\n",
    "            plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot mutual information between Z and S"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams.update({'font.size': 16})\n",
    "plt.rcParams['figure.dpi'] = 300\n",
    "linestyles = ['-', '--', '-.', ':']\n",
    "colors = ['b', 'r', 'g', 'c']\n",
    "for i in range(mi_s_ts.shape[0]):\n",
    "    plt.figure(figsize=(15,8))\n",
    "#     plt.title(r'Mutual Information with $\\gamma$ $\\in$ [0.0001, 0.1, 10, 100, 1000]' + f' dz_{DIM_Z[i]} - {exp_info}', fontsize=20)\n",
    "    for j in range(mi_s_ts.shape[1]):\n",
    "        plt.plot(mi_s_ts[i][j], label=r'Test, $d_z=$' + f'{DIM_Z[i]}' + r',$\\lambda=$' + f\"{lam_list[j]}\", linestyle=linestyles[0], c=colors[j%len(colors)] ,linewidth=3)\n",
    "#         plt.plot(mi_s_tr[i][j], label=r'Train, $d_z=$' + f'{DIM_Z[i]}' + r',$\\lambda=$' + f\"{lam_list[j]}\", linestyle=linestyles[1], c=colors[j%len(colors)] ,linewidth=3)\n",
    "\n",
    "    # plt.xscale('linear')\n",
    "    plt.xlabel(r'$\\gamma$', fontsize=24)\n",
    "    plt.xticks(list(range(len(gamma_list))), gamma_list, rotation=90)\n",
    "\n",
    "    plt.ylabel(r'I($\\mathbf{S}$;$\\mathbf{Z}$)', fontsize=24)\n",
    "    plt.legend(prop={'size': 16})\n",
    "    plt.grid()\n",
    "    plt.savefig(f'./saved_figures/chart_colored_mnist_mi_zs_d_{DIM_Z[i]}_{exp_info}_lam_.svg', format='svg', bbox_inches='tight')\n",
    "    plt.savefig(f'./saved_figures/chart_colored_mnist_mi_zs_d_{DIM_Z[i]}_{exp_info}_lam_.png', format='png', bbox_inches='tight')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot mutual information between Z and U"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if supervised:\n",
    "    plt.rcParams.update({'font.size': 16})\n",
    "    plt.rcParams['figure.dpi'] = 300\n",
    "    plt.figure(figsize=(15,8))\n",
    "#     plt.title(r'Mutual Information with $\\gamma$ $\\in$ [0.0001, 0.1, 10, 100, 1000]' + f' - {exp_info}', fontsize=20)\n",
    "    linestyles = ['-', '--', '-.', ':']\n",
    "    colors = ['b', 'r', 'g', 'c']\n",
    "    for i in range(mi_u_ts.shape[0]):\n",
    "        for j in range(mi_u_ts.shape[1]):\n",
    "            plt.plot(mi_u_ts[i][j], label=r'Test, $d_z=$' + f'{DIM_Z[i]}' + r',$\\lambda=$' + f\"{lam_list[j]}\", linestyle=linestyles[0], c=colors[j%len(colors)] ,linewidth=3)\n",
    "    #         plt.plot(mi_u_tr[i][j], label=r'Train, $d_z=$' + f'{DIM_Z[i]}' + r',$\\lambda=$' + f\"{lam_list[j]}\", linestyle=linestyles[1], c=colors[j%len(colors)] ,linewidth=3)\n",
    "\n",
    "    # plt.xscale('linear')\n",
    "    plt.xlabel(r'$\\gamma$', fontsize=24)\n",
    "    plt.xticks(list(range(len(gamma_list))), gamma_list, rotation=90)\n",
    "\n",
    "    plt.ylabel(r'I($\\mathbf{U}$;$\\mathbf{Z}$)', fontsize=24)\n",
    "    plt.legend(prop={'size': 16})\n",
    "    plt.grid()\n",
    "    plt.savefig(f'./saved_figures/chart_colored_mnist_mi_zu_{exp_info}.svg', format='svg', bbox_inches='tight')\n",
    "    plt.savefig(f'./saved_figures/chart_colored_mnist_mi_zu_{exp_info}.png', format='png', bbox_inches='tight')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, u_test_hat = autoencoder.predict(x_test)\n",
    "if supervised:\n",
    "    print(\"Acc = \", np.mean(np.argmax(u_test_hat,axis=1)==np.argmax(u_test,axis=1)) * 100)\n",
    "else:\n",
    "    print(\"MSE = \", print(np.mean(np.sum(np.square(u_train_hat-u_train),axis=(1,2,3)))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not supervised:\n",
    "    def save_figures(pib):\n",
    "        n_plot_samps = 10\n",
    "        ind_list = [3, 2, 1, 18, 4 ,15, 11, 0, 128, 9]\n",
    "\n",
    "        _, x_hat = pib.predict(x_test[ind_list])\n",
    "\n",
    "        from utils.pics_tools import plot_image_grid\n",
    "        fig = plot_image_grid(np.concatenate([x_test[ind_list], x_hat], axis=0),[IMG_DIM, IMG_DIM, 3],(2,n_plot_samps))\n",
    "        fig.set_size_inches(10, 2.1)\n",
    "        plt.subplots_adjust(wspace=0, hspace=0)\n",
    "    #     plt.savefig(f'./saved_figures/colored_mnist_unblanced_{exp_info}.svg', format='svg')\n",
    "    #     plt.savefig(f'./saved_figures/colored_mnist_unblanced_{exp_info}.png', format='png')\n",
    "        plt.show()\n",
    "    save_figures(autoencoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Save results values to file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.io as sio\n",
    "sio.savemat(f'./saved_data/acc_colored_mnist_{exp_info}.mat', {'util_acc_ts':util_acc_ts, 'util_acc_tr':util_acc_tr, 'sens_acc_ts':sens_acc_ts, 'sens_acc_tr':sens_acc_tr})\n",
    "sio.savemat(f'./saved_data/mi_colored_mnist_{exp_info}.mat', {'mi_s_ts':mi_s_ts, 'mi_s_tr':mi_s_tr, 'mi_u_ts':mi_u_ts, 'mi_u_tr':mi_u_tr})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load results from file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.io as sio\n",
    "mat_contents  = sio.loadmat(f'./saved_data/acc_colored_mnist_{exp_info}.mat')\n",
    "util_acc_ts = mat_contents['util_acc_ts']\n",
    "util_acc_tr = mat_contents['util_acc_tr']\n",
    "sens_acc_ts = mat_contents['sens_acc_ts']\n",
    "sens_acc_tr = mat_contents['sens_acc_tr']\n",
    "\n",
    "\n",
    "mat_contents  = sio.loadmat(f'./saved_data/mi_colored_mnist_{exp_info}.mat')\n",
    "mi_s_tr = mat_contents['mi_s_tr']\n",
    "mi_s_ts = mat_contents['mi_s_ts']\n",
    "mi_u_ts = mat_contents['mi_u_ts']\n",
    "mi_u_tr = mat_contents['mi_u_tr']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lam_list = [0.0001, 0.1, 10, 100, 1000]\n",
    "gamma_list = [0.0001, 0.1, 10, 100, 1000]\n",
    "DIM_Z = [8, 64]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:tensorflow-24]",
   "language": "python",
   "name": "conda-env-tensorflow-24-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}