WGAN.py

# https://machinelearningmastery.com/how-to-code-a-wasserstein-generative-adversarial-network-wgan-from-scratch/
# -*- coding: utf-8 -*-
"""Shakti__telecom_wcgan-gp.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1ZoMl4TM8Ell_2ClnJAifPEr-aCdBHVUg
"""

# !pip install table-evaluator

# from table_evaluator import TableEvaluator

import tensorflow as tf
import keras
from functools import partial
import itertools
import math
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from keras.layers import Input, Dense, Flatten, Dropout, Embedding, multiply, LeakyReLU
from keras.models import Sequential, Model
# from keras.optimizers.legacy import Adam
from keras.optimizers import Adam  
import keras.backend as K
from tensorflow.python.framework.ops import disable_eager_execution
disable_eager_execution()
import warnings
warnings.simplefilter("ignore")

# import random as python_random
# from numpy.random import randn
# from tensorflow import random

import random
random.seed(101)
seed = 101 #108 # to make results replicable (much better than 102, 103)
np.random.seed(seed) # for numpy
tf.random.set_seed(seed) # for tensorflow
keras.utils.set_random_seed(seed) # keras
# python_random.seed(seed) # for python
optimizers = Adam(learning_rate=0.0004, beta_1=0.05, beta_2=0.9)
#Check how many GPUs there is avaliable for training:

# import tensorflow as tf

# gpus = tf.config.list_physical_devices('GPU')
# logical_gpus = tf.config.experimental.list_logical_devices('GPU')
# print("Number of GPUs & Logical GPUs available: ", len(gpus), len(logical_gpus))

#--- [1] Read and transform data

import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler, Normalizer, LabelEncoder
from sklearn.model_selection import train_test_split
from sklearn.decomposition import PCA

df = pd.read_csv("https://raw.githubusercontent.com/shakti2594/Shared_dataset/main/Telco-Customer-Churn.csv")

# keep minority group only (Churn = 'Yes')
df.drop(df[(df['Churn'] == 'No')].index, inplace=True)

#- [1.1] transforming TotalCharges to TotalChargeResidues, add to dataframe

df['TotalCharges'] = pd.to_numeric(df['TotalCharges'], errors='coerce')
print(df.isna().sum().sum())
df = df.dropna()
print(df.shape)
arr1 = df['tenure'].to_numpy()
arr2 = df['TotalCharges'].to_numpy()
arr2 = arr2.astype(float)
residues = arr2 - arr1 * np.sum(arr2) / np.sum(arr1) # also try arr2/arr1
df['TotalChargeResidues'] = residues
df = df[['tenure', 'MonthlyCharges', 'TotalChargeResidues', 'Churn']]

#- [1.2] adding labels to data, to boost GAN

le = LabelEncoder()
df['Churn'] = le.fit_transform(df['Churn'])
encoded_labels = {num:label for (num, label) in zip(range(15), le.classes_)}
encoded_labels

x_train, x_test, y_train, y_test = train_test_split(df.iloc[:,:-1],
                                                    df['Churn'],
                                                    test_size=1/7.0,
                                                    random_state=0)
x_train = np.array(x_train)
x_test = np.array(x_test)
#y_train = y_train.values.reshape(-1,1)
#y_test = y_test.values.reshape(-1,1)
y_train = np.array(y_train)
y_test = np.array(y_test)

#- [1.3] potential data transforms

ss = StandardScaler().fit(x_train)
x_train = ss.transform(x_train)
x_test = ss.transform(x_test)
#print(x_train.shape, x_test.shape)

pca = PCA(.99).fit(x_train)
x_train = pca.transform(x_train)
x_test = pca.transform(x_test)
#print(x_train.shape, x_test.shape)

norm = Normalizer().fit(x_train)
x_train = norm.transform(x_train)
x_test = norm.transform(x_test)
#print(x_train.shape, x_test.shape)

print(x_train.shape, x_test.shape)
print(x_train)


#--- [2] Model Definition"""

class RandomWeightedAverage(tf.keras.layers.Layer):
    # Provides a (random) weighted average between real and generated samples

    def __init__(self, batch_size):
        super().__init__()
        self.batch_size = batch_size
        

    def call(self, inputs, **kwargs):
        alpha = tf.random.uniform((self.batch_size, 1))
        return (alpha * inputs[0]) + ((1 - alpha) * inputs[1])

    def compute_output_shape(self, input_shape):
        return input_shape[0]

class WCGANGP():

    def __init__(self,
                 x_train,
                 y_train,
                 latent_dim: int,
                 batch_size: int,
                 n_critic: int):
        """Implement WCGAN with Gradient Penalty."""
        
        self.x_train = x_train
        self.y_train = y_train

        self.num_classes = len(np.unique(y_train))
        self.data_dim = x_train.shape[1]

        self.latent_dim = latent_dim
        self.batch_size = batch_size

        self.n_critic = n_critic

        # Log training progress.
        self.losslog = []

        # Adam optimizer, suggested by original paper.
        optimizer = optimizers

        # Build the generator and critic
        self.generator = self.build_generator()
        self.critic = self.build_critic()

        # Freeze generator's layers while training critic.
        self.generator.trainable = False

        # Data input (real sample).
        real_data = Input(shape=self.data_dim)
        # Noise input (z).
        noise = Input(shape=(self.latent_dim,))
        # Label input.
        label = Input(shape=(1,))

        # Generate data based of noise (fake sample)
        fake_data = self.generator([noise, label])

        # Critic (discriminator) determines validity of the real and fake images.
        fake = self.critic([fake_data, label])
        valid = self.critic([real_data, label])

        # Construct weighted average between real and fake images.
        interpolated_data = RandomWeightedAverage(self.batch_size)([real_data, fake_data])

        # Determine validity of weighted sample.
        validity_interpolated = self.critic([interpolated_data, label])

        # Use Python partial to provide loss function with additional
        # 'averaged_samples' argument.
        partial_gp_loss = partial(self.gradient_penalty_loss,
                          averaged_samples=interpolated_data)
        # Keras requires function names.
        partial_gp_loss.__name__ = 'gradient_penalty'

        self.critic_model = Model(inputs=[real_data, label, noise],
                            outputs=[valid, fake, validity_interpolated])

        self.critic_model.compile(loss=[self.wasserstein_loss,
                                        self.wasserstein_loss,
                                        partial_gp_loss],
                                  optimizer=optimizer,
                                  loss_weights=[1, 1, 10])

        # For the generator we freeze the critic's layers.
        self.critic.trainable = False
        self.generator.trainable = True

        # Sampled noise for input to generator.
        noise = Input(shape=(self.latent_dim,))

        # Add label to input.
        label = Input(shape=(1,))

        # Generate data based of noise.
        fake_data = self.generator([noise, label])
        

        # Discriminator determines validity.
        valid = self.critic([fake_data, label])

        # Define generator model.
        self.generator_model = Model([noise, label], valid)
        self.generator_model.compile(loss=self.wasserstein_loss,
                                     optimizer=optimizer)

    def gradient_penalty_loss(self, y_true, y_pred, averaged_samples):
     
        ## Computes gradient penalty based on prediction and weighted real / fake samples
    
        gradients = K.gradients(y_pred, averaged_samples)[0]

        # compute the euclidean norm by squaring ...
        gradients_sqr = K.square(gradients)
        #   ... summing over the rows ...
        gradients_sqr_sum = K.sum(gradients_sqr,
                                  axis=np.arange(1, len(gradients_sqr.shape)))
        #   ... and sqrt
        gradient_l2_norm = K.sqrt(gradients_sqr_sum)

        # compute lambda * (1 - ||grad||)^2 still for each single sample
        gradient_penalty = K.square(1 - abs(gradient_l2_norm))

        # return the mean as loss over all the batch samples
        return K.mean(gradient_penalty)


    def wasserstein_loss(self, y_true, y_pred):
        # Computes Wasserstein loss from real and fake predictions.
        return K.mean(y_true * y_pred)

    def build_generator(self):

        model = Sequential(name="Generator")

        # First hidden layer.
        model.add(Dense(256, input_dim=self.latent_dim))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dropout(0.3))

        # Second hidden layer.
        model.add(Dense(512))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dropout(0.3))

        # Third hidden layer.
        model.add(Dense(1024))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dropout(0.3))

        # Output layer.
        model.add(Dense(self.data_dim, activation="tanh"))

        model.summary()

        # Noise and label input layers.
        noise = Input(shape=(self.latent_dim,))
        label = Input(shape=(1,), dtype="int32")

        # Embed labels into onehot encoded vectors.
        label_embedding = Flatten()(Embedding(self.num_classes, self.latent_dim)(label))

        # Multiply noise and embedded labels to be used as model input.
        model_input = multiply([noise, label_embedding])

        generated_data = model(model_input)
        return Model([noise, label], generated_data, name="Generator")

    def build_critic(self):

        model = Sequential(name="Critic")

        # First hidden layer.
        model.add(Dense(1024, input_dim=self.data_dim))
        model.add(LeakyReLU(alpha=0.2))

        # Second hidden layer.
        model.add(Dense(512))
        model.add(LeakyReLU(alpha=0.2))

        # Third hidden layer.
        model.add(Dense(256))
        model.add(LeakyReLU(alpha=0.2))

        # Output layer with linear activation.
        model.add(Dense(1))
        model.summary()

        # Artificial data input.
        generated_sample = Input(shape=self.data_dim)
        # Label input.
        label = Input(shape=(1,), dtype="int32")

        # Embedd label as onehot vector.
        label_embedding = Flatten()(Embedding(self.num_classes, self.data_dim)(label))

        # Multiply fake data sample with label embedding to get critic input.
        model_input = multiply([generated_sample, label_embedding])

        validity = model(model_input)
        return Model([generated_sample, label], validity, name="Critic")


#--- [3] Main part: training GAN (epoch iteration)

    def train(self, epochs):

        # Adversarial ground truths.
        valid = -(np.ones((self.batch_size, 1)))
        fake =  np.ones((self.batch_size, 1))
        dummy = np.zeros((self.batch_size, 1))

        # Number of batches.
        self.n_batches = math.floor(self.x_train.shape[0] / self.batch_size)
        overhead = self.x_train.shape[0] % self.batch_size

        for epoch in range(epochs):

            # Reset training set.
            self.x_train = x_train.copy()
            self.y_train = y_train.copy()

            # Select random overhead rows that do not fit into batches.
            rand_overhead_idx = np.random.choice(range(self.x_train.shape[0]), overhead, replace=False)

            # Remove random overhead rows.
            self.x_train = np.delete(self.x_train, rand_overhead_idx, axis=0)
            self.y_train = np.delete(self.y_train, rand_overhead_idx, axis=0)

            # Split training data into batches.
            x_batches = np.split(self.x_train, self.n_batches)
            y_batches = np.split(self.y_train, self.n_batches)


            for x_batch, y_batch, i in zip(x_batches, y_batches, range(self.n_batches)):

                for _ in range(self.n_critic):

                    # ---------------------
                    #  Train Critic
                    # ---------------------

                    # Generate random noise.
                    noise = np.random.normal(0, 1, (self.batch_size, self.latent_dim))

                    # Train the critic.
                    d_loss = self.critic_model.train_on_batch(
                        [x_batch, y_batch, noise],
                        [valid, fake, dummy])


                # ---------------------
                #  Train Generator
                # ---------------------

                # Generate sample of artificial labels.
                generated_labels = np.random.randint(0, self.num_classes, self.batch_size).reshape(-1, 1)

                # Train generator.
                g_loss = self.generator_model.train_on_batch([noise, generated_labels], valid)


                # ---------------------
                #  Logging
                # ---------------------

                self.losslog.append([d_loss[0], g_loss])

                DLOSS = "%.4f" % d_loss[0]
                GLOSS = "%.4f" % g_loss

                if i % 100 == 0:
                    print(f"{epoch} - {i}/{self.n_batches} \t [D loss: {DLOSS}] [G loss: {GLOSS}]")

#- Model training: core of the algorithm

wcgan = WCGANGP(
    x_train,
    y_train,
    latent_dim=32,
    batch_size=128,
    n_critic=5)

n_epochs = 500
wcgan.train(epochs=n_epochs)


#--- [4] Synthetic data generation 

label_ratios = { label:len(y_train[y_train == label])/y_train.shape[0] for label in np.unique(y_train) }
label_ratios

def generate_samples(n: int, latent_dim: int):

    """Use WCGAN to generate new data."""
    noise = np.random.normal(0, 1, (n, latent_dim))

    # Create sampled labels.
    sampled_labels = [
        np.full(round(ratio*n), label).tolist()
        for label, ratio in label_ratios.items()
    ]

    # Convert list to numpy array.
    sampled_labels = np.array((list(itertools.chain(*sampled_labels))))

    # Use CGAN to generate aritficial data.
    return wcgan.generator.predict([noise, sampled_labels])

generated_samples = generate_samples(x_train.shape[0], wcgan.latent_dim)

# real_samples = pd.DataFrame(x_train)
df_real = pd.DataFrame(x_train, columns = ['tenure', 'MonthlyCharges', 'TotalChargeResidues'])
df_synth = pd.DataFrame(generated_samples, columns = ['tenure', 'MonthlyCharges', 'TotalChargeResidues'])
df_real.to_csv('WGAN_telecom_real.csv') 
df_synth.to_csv('WGAN_telecom_synth.csv') 
print(df_synth)



#--- [5] Synthetic data evaluation

# table_evaluator = TableEvaluator(
#     real_samples,
#     pd.DataFrame(generated_samples)
# )

# table_evaluator.visual_evaluation()

plt.figure(figsize=(10,5))
#plt.ylim(-6, 3)
plt.plot(wcgan.losslog)
plt.title("WCGANGP Losses")
plt.xlabel("Batches")
plt.ylabel("Loss")
plt.legend(['Critic loss', 'Generator loss'])
plt.show()
"""
plt.figure(figsize=(10,5))
# plt.ylim(-6, 2)
plt.plot(wcgan.losslog)
plt.title("WCGANGP Losses")
plt.xlabel("Batches")
plt.ylabel("Loss")
plt.legend(['Critic loss', 'Generator loss'])

plt.figure(figsize=(10,5))
# plt.ylim(-6, 2)
plt.plot(wcgan.losslog)
plt.title("WCGANGP Losses")
plt.xlabel("Batches")
plt.ylabel("Loss")
plt.legend(['Critic loss', 'Generator loss'])

plt.figure(figsize=(10,5))
# plt.ylim(-6, 2)
plt.plot(wcgan.losslog)
plt.title("WCGANGP Losses")
plt.xlabel("Batches")
plt.ylabel("Loss")
plt.legend(['Critic loss', 'Generator loss'])

plt.figure(figsize=(10,5))
plt.ylim(-6, 2)
plt.plot(wcgan.losslog)
plt.title("WCGANGP Losses")
plt.xlabel("Batches")
plt.ylabel("Loss")
plt.legend(['Critic loss', 'Generator loss'])
"""
"""**Saving models**"""

# wcgan.generator.save("models/generator.h5")
# wcgan.critic.save("models/critic.h5")