nevoit/Synthesizing-Tabular-Data-Using-GAN

1import os2import pickle3import matplotlib.pyplot as plt4import numpy as np5from keras.layers import Dense, Dropout, LeakyReLU6from keras.models import Sequential7from keras.optimizers import Adam8from numpy.random import randn9from sklearn.ensemble import RandomForestClassifier10from sklearn import metrics11from tqdm import tqdm12 13 14class GG(object):15 16    def __init__(self, number_of_features, saved_models_path, learning_rate, dropout, alpha):17        """18        The constructor for the General Generator class.19        :param number_of_features: Number of features in the data. Used to determine the noise dimensions20        :param saved_models_path: The folder where we save the models.21        """22        self.saved_models_path = saved_models_path23        self.number_of_features = number_of_features24 25        self.generator_model = None26        self.discriminator_model = RandomForestClassifier()27        self.dropout = dropout28        self.alpha = alpha29        self.noise_dim = int(number_of_features / 2)30        self.learning_rate = learning_rate31        self.number_of_features = number_of_features32        self.build_generator()  # build the generator.33        self.losses = {'gen_loss': [], 'dis_loss_pred': [], 'dis_loss_proba': []}34        # self.results = {}35 36    def build_generator(self):37        """38        This function creates the generator model for the GG.39        We used a fairly simple MLP architecture.40        :return:41        """42 43        self.generator_model = Sequential()44        self.generator_model.add(Dense(int(self.number_of_features * 2), input_shape=(self.noise_dim + 1, )))45        self.generator_model.add(LeakyReLU(alpha=self.alpha))46 47        self.generator_model.add(Dense(int(self.number_of_features * 4)))48        self.generator_model.add(LeakyReLU(alpha=self.alpha))49        self.generator_model.add(Dropout(self.dropout))50 51        self.generator_model.add(Dense(int(self.number_of_features * 2)))52        self.generator_model.add(LeakyReLU(alpha=self.alpha))53        self.generator_model.add(Dropout(self.dropout))54 55        self.generator_model.add(Dense(self.number_of_features, activation='sigmoid'))56        optimizer = Adam(lr=self.learning_rate)57        self.generator_model.compile(loss='categorical_crossentropy', optimizer=optimizer)58 59        # self.generator_model.summary()60 61    def train_gg(self, x_train, y_train, epochs, batch_size, model_name, data, output_path, to_plot=False):62        """63        This function running the training stage manually.64        :param output_path: Path to save loss fig65        :param to_plot: Plots the losses if True66        :param x_train: the training set features67        :param y_train: the training set classes68        :param model_name: name of model to save (for generator)69        :param epochs: number of epochs70        :param batch_size: the batch size71        :return: trains the discriminator and generator.72        """73 74        losses_path = os.path.join(self.saved_models_path, f'{model_name}_losses')75        model_file = os.path.join(self.saved_models_path, f'{model_name}_part_2_gen_weights.h5')76        77        # First train the discriminator78        self.train_black_box_dis(x_train, y_train)79        self.train_generator(x_train, model_file, epochs, batch_size, losses_path)80        if to_plot:81            self.plot_losses(data, output_path)82 83    def train_black_box_dis(self, x_train, y_train):84        """85        Trains the discriminator and saves it.86        :param x_train: the training set features87        :param y_train: the training set classes88        :return:89        """90        dis_output = os.path.join(self.saved_models_path, 'black_box_dis_model')91 92        if os.path.exists(dis_output):93            # print('Blackbox discriminator already trained')94            with open(dis_output, 'rb') as rf_file:95                self.discriminator_model = pickle.load(rf_file)96 97        self.discriminator_model.fit(x_train, y_train)98        with open(dis_output, 'wb') as rf_file:99            pickle.dump(self.discriminator_model, rf_file)100 101    def train_generator(self, data, model_path, epochs, start_batch_size, losses_path):102        """103        Function for training the general generator.104        :param losses_path: The filepath for the loss results105        :param data: The normalized dataset106        :param model_path: The name of the model to save. includes epoch size, batches etc.107        :param epochs: Number of epochs108        :param start_batch_size: Size of batch to use.109        :return: trains the generator, saves it and the losses during training.110        """111 112        if os.path.exists(model_path):113            self.generator_model.load_weights(model_path)114            with open(losses_path, 'rb') as loss_file:115                self.losses = pickle.load(loss_file)116            return117 118        for epoch in range(epochs):  # iterates over the epochs119            np.random.shuffle(data)120            batch_size = start_batch_size121            for i in tqdm(range(0, data.shape[0], batch_size), ascii=True):  # Iterate over batches122                if data.shape[0] - i >= batch_size:123                    batch_input = data[i:i + batch_size]124                else:  # The last iteration125                    batch_input = data[i:]126                    batch_size = batch_input.shape[0]127 128                g_loss = self.train_generator_on_batch(batch_input)129                self.losses['gen_loss'].append(g_loss)130 131        self.save_generator_model(model_path, losses_path)132 133    def save_generator_model(self, generator_model_path, losses_path):134        """135        Saves the model and the loss data with pickle.136 137        :param generator_model_path: File path for the generator138        :param losses_path: File path for the losses139        :return:140        """141        self.generator_model.save_weights(generator_model_path)142        with open(losses_path, 'wb+') as loss_file:143            pickle.dump(self.losses, loss_file)144 145    def train_generator_on_batch(self, batch_input):146        """147        Trains the generator for a single batch. Creates the necessary input, comprised of noise and the real148        probabilities obtained from the black box. Compared to the target output, made of real samples and the149        probabilities made up by the generator.150        :param batch_input:151        :return:152        """153        batch_size = batch_input.shape[0]154        discriminator_probabilities = self.discriminator_model.predict_proba(batch_input)[:, -1:]155        # noise = randn(self.noise_dim * batch_size).reshape((batch_size, self.noise_dim))156 157        noise = randn(batch_size, self.noise_dim)158        gen_model_input = np.hstack([noise, discriminator_probabilities])159        generated_probabilities = self.generator_model.predict(gen_model_input)[:, -1:]  # Take only probabilities160        target_output = np.hstack([batch_input, generated_probabilities])161        g_loss = self.generator_model.train_on_batch(gen_model_input, target_output)  # The actual training162 163        return g_loss164 165    def plot_discriminator_results(self, x_test, y_test, data, path):166        """167        :param x_test: Test set168        :param y_test: Test classes169        :return: Prints the required plots.170        """171 172        blackbox_probs = self.discriminator_model.predict_proba(x_test)173        discriminator_predictions = self.discriminator_model.predict(x_test)174        count_1 = int(np.sum(y_test))175        count_0 = int(y_test.shape[0] - count_1)176        class_data = (['Class 0', 'Class 1'], [count_0, count_1])177        self.plot_data(class_data, path, mode='bar', x_title='Class', title=f'Distribution of classes - {data} dataset')178        self.plot_data(blackbox_probs[:, 0], path, title=f'Probabilities for test set - class 0 - {data} dataset')179        self.plot_data(blackbox_probs[:, 1], path, title=f'Probabilities for test set - class 1 - {data} dataset')180 181        min_confidence = blackbox_probs[:, 0].min(), blackbox_probs[:, 1].min()182        max_confidence = blackbox_probs[:, 0].max(), blackbox_probs[:, 1].max()183        mean_confidence = blackbox_probs[:, 0].mean(), blackbox_probs[:, 1].mean()184 185        print("Accuracy:", metrics.accuracy_score(y_test, discriminator_predictions))186        for c in [0, 1]:187            print(f'Class {c} - Min confidence: {min_confidence[c]} - Max Confidence: {max_confidence[c]} - '188                  f'Mean confidence: {mean_confidence[c]}')189 190    def plot_generator_results(self, data, path, num_of_instances=1000):191        """192        Creates plots for the generator results on 1000 instances.193        :param path:194        :param data: Name of dataset used.195        :param num_of_instances: Number of samples to generate.196        :return:197        """198        sampled_proba, generated_instances = self.generate_n_samples(num_of_instances)199 200        proba_fake = self.discriminator_model.predict_proba(generated_instances[:, :-1])201        for c in [0, 1]:202            title = f'Confidence Score for Class {c} of Fake Samples - {data} dataset'203            self.plot_data(proba_fake[:, c], path, x_title='Confidence Score', title=title)204 205        black_box_confidence = proba_fake[:, 1:]206        proba_error = np.abs(sampled_proba - black_box_confidence)207        generated_classes = np.array([int(round(c)) for c in generated_instances[:, -1].tolist()]).reshape(1000, 1)208        proba_stats = np.hstack([sampled_proba, generated_classes, proba_fake[:, :1], proba_fake[:, 1:], proba_error])209 210        for c in [0, 1]:211            class_data = proba_stats[proba_stats[:, 1] == c]212            class_data = class_data[class_data[:, 0].argsort()]  # Sort it for the plot213            title = f'Error rate for different probabilities, class {c} - {data} dataset'214            self.plot_data((class_data[:, 0], class_data[:, -1]), path, mode='plot', y_title='error rate', title=title)215 216    def generate_n_samples(self, n):217        """218        Functions for generating N samples with a uniformly distribution confidence level.219        :param n: Number of samples220        :return: a tuple of the confidence scores used and the samples created.221        """222        noise = randn(n, self.noise_dim)223        # confidences = np.sort(np.random.uniform(0, 1, (n, 1)), axis=0)224        confidences = np.random.uniform(0, 1, (n, 1))225 226        generator_input = np.hstack([noise, confidences])  # Stick them together227        generated_instances = self.generator_model.predict(generator_input)  # Create samples228 229        return confidences, generated_instances230 231    @staticmethod232    def plot_data(data, path, mode='hist', x_title='Probabilities', y_title='# of Instances', title='Distribution'):233        """234        :param path: Path to save235        :param mode: Mode to use236        :param y_title: Title of y axis237        :param x_title: Title of x axis238        :param data: Data to plot239        :param title: Title of plot240        :return: Prints a plot241        """242        plt.clf()243 244        if mode == 'hist':245            plt.hist(data)246        elif mode == 'bar':247            plt.bar(data[0], data[1])248        else:249            plt.plot(data[0], data[1])250 251        plt.title(title)252        plt.ylabel(y_title)253        plt.xlabel(x_title)254        # plt.show()255        path = os.path.join(path, title)256        plt.savefig(path)257 258    def plot_losses(self, data, path):259        """260        Plot the losses while training261        :return:262        """263        plt.clf()264        plt.plot(self.losses['gen_loss'])265        plt.title('Model loss')266        plt.ylabel('Loss')267        plt.xlabel('Iteration')268        # plt.show()269        plt.savefig(os.path.join(path, f'{data} dataset - general_generator_loss.png'))270 271    def get_error(self, num_of_instances=1000):272        """273        Calculates the error of the generator we created by measuring the difference between the probability that274        was given as input and the probability of the discriminator on the sample created.275        :param num_of_instances: Number of samples to generate.276        :return: An array of errors.277        """278        sampled_proba, generated_instances = self.generate_n_samples(num_of_instances)279        proba_fake = self.discriminator_model.predict_proba(generated_instances[:, :-1])280        black_box_confidence = proba_fake[:, 1:]281        return np.abs(sampled_proba - black_box_confidence)282 283