基于keras 的AdversarialAutoencoder实例

本次学习目标为实现非监督AdversarialAutoencoder。

以下部分为可能用得到的package

from keras.layers import Dense, Flatten, Reshape, Conv2D, UpSampling2D, Cropping2D, MaxPool2D
from keras.models import Sequential
from keras.layers import Input
from keras.models import Model
import numpy as np
import tensorflow as tf
import datetime
from tensorflow.keras.layers.experimental import preprocessing
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

discriminator 初始化：

def build_custom_discriminator(latent_dim, initializer, info=False):
    discriminator = Sequential()
    discriminator.add(Dense(100, activation='relu',
                            kernel_initializer=initializer,
                            input_dim=latent_dim, bias_initializer=initializer)
                      )
    discriminator.add(Dense(500, activation='relu',
                            kernel_initializer=initializer,
                            bias_initializer=initializer))
    discriminator.add(Dense(1, activation="sigmoid",
                            kernel_initializer=initializer,
                            bias_initializer=initializer))

    if info:
        print(discriminator.summary())

    return discriminator

encoder 与 decoder的初始化（在这部分我使用了卷积conv2d方法）：

def build_denser_omata_encoder_decoder(input_shape, latent_dim, initializer, info=False, act="elu", dense_act=None):

    encoder = Sequential()
    encoder.add(Conv2D(32, (5, 5), padding="same", activation=act, input_shape=input_shape, kernel_initializer=initializer))
    encoder.add(MaxPool2D(padding="same"))
    encoder.add(Conv2D(64, (3, 3), activation=act, padding="same", kernel_initializer=initializer))
    encoder.add(MaxPool2D(padding="same"))
    encoder.add(Conv2D(128, (3, 3), activation=act, padding="same", kernel_initializer=initializer))
    encoder.add(MaxPool2D(padding="same"))
    encoder.add(Flatten())
    encoder.add(Dense(int(5376/2),
                      kernel_initializer=initializer,
                      activation=dense_act))
    encoder.add(Dense(latent_dim, activation="linear"))

    if info:
        print(encoder.summary())

    decoder = Sequential()
    decoder.add(Dense(int(5376/2),
                      kernel_initializer=initializer,
                      activation=dense_act,
                      input_shape=(latent_dim,)))
    decoder.add(Dense(5376,
                      kernel_initializer=initializer,
                      activation=dense_act,
                      input_shape=(int(5376/2),)))
    decoder.add(Reshape((encoder.layers[6].input_shape[1],
                         encoder.layers[6].input_shape[2], 128)))
    decoder.add(Conv2D(128, (3, 3), activation=act, padding="same",
                       kernel_initializer=initializer))
    decoder.add(UpSampling2D())
    decoder.add(Conv2D(64, (3, 3), activation=act, padding="same",
                       kernel_initializer=initializer))
    decoder.add(UpSampling2D())
    decoder.add(Conv2D(32, (3, 3), activation=act, padding="same",
                       kernel_initializer=initializer))
    decoder.add(UpSampling2D())
    decoder.add(Conv2D(2, (3, 3), activation="linear", padding="same",
                kernel_initializer=initializer))
    decoder.add(Cropping2D(cropping=((1, 0), (3, 3))))

    decoder.build(input_shape)

    if info:
        print(decoder.summary())

    return encoder, decoder

初始化与训练：

class AAE:

    def __init__(self, encoder, decoder, discriminator, optimizer):
        self.encoder = encoder
        self.decoder = decoder
        self.discriminator = discriminator
        self.latent_dim = self.decoder.layers[0].input_shape[1]

        self.optimizer = optimizer

    def compile(self, input_shape):
        """
        Compilation of models according to original paper on adversarial
        autoencoders
        Args:
            input_shape (tuple): Shape of input data
        """

        self.input_shape = input_shape

        grid = Input(shape=self.input_shape)
        encoded_repr = self.encoder(grid)
        gen_grid = self.decoder(encoded_repr)
        self.autoencoder = Model(grid, gen_grid)

        valid = self.discriminator(encoded_repr)
        self.encoder_discriminator = Model(grid, valid)

        self.discriminator.compile(optimizer=self.optimizer,loss='binary_crossentropy',metrics=['accuracy'])

        self.autoencoder.compile(optimizer=self.optimizer,loss='mse',metrics=['accuracy'])

        self.discriminator.trainable = False

        self.encoder_discriminator.compile(optimizer=self.optimizer,loss='binary_crossentropy',metrics=['accuracy'])

    def train(self, train_data, epochs, val_data=None, batch_size=128,
              val_batch_size=128, wandb_log=False):
        """
        Training model according to original paper on adversarial autoencoders
        Args:
            train_data (np.ndarray): Array with train data
            epochs (int): Number of epochs
            val_data (np.ndarray, optional): Array with validation data.
                Defaults to None.
            batch_size (int, optional): Batch size. Defaults to 128.
            plot_losses (bool, optional): Whether to plot losses.
                Defaults to False.
            print_losses (bool, optional): Whether to print losses.
                Defaults to False.
        """

        d_loss_val = g_loss_val = None

        train_dataset = tf.data.Dataset.from_tensor_slices(train_data)
        train_dataset = train_dataset.shuffle(buffer_size=train_data.shape[0],reshuffle_each_iteration=True).batch(batch_size, drop_remainder=True)

        if val_data is not None:
            val_dataset = tf.data.Dataset.from_tensor_slices(val_data)
            val_dataset = val_dataset.shuffle(
                buffer_size=val_data.shape[0],
                reshuffle_each_iteration=True).\
                batch(val_batch_size, drop_remainder=True)

        # Set up tensorboard logging
        current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
        train_log_dir = 'logs/' + current_time + '/train'
        val_log_dir = 'logs/' + current_time + '/val'
        train_summary_writer = tf.summary.create_file_writer(train_log_dir)
        val_summary_writer = tf.summary.create_file_writer(val_log_dir)

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

        for epoch in range(epochs):

            # Reconstruction phase
            loss_cum = 0
            acc_cum = 0
            for step, grids in enumerate(train_dataset):
                # Train the autoencoder reconstruction
                loss, acc = self.autoencoder.train_on_batch(grids, grids)
                loss_cum += loss
                acc_cum += acc

            # Average the loss and accuracy over the entire dataset
            loss = loss_cum/(step+1)
            acc = acc_cum/(step+1)

            # Regularization phase
            d_loss_cum = 0
            g_loss_cum = 0
            for step, grids in enumerate(train_dataset):

                # Generate real and fake latent space. Fake latent space is
                # the normal distribution
                latent_fake = self.encoder.predict(grids)
                latent_real = np.random.normal(size=(batch_size,self.latent_dim))

                # Train the discriminator
                d_loss_real = self.discriminator.train_on_batch(latent_real,valid)[0]
                d_loss_fake = self.discriminator.train_on_batch(latent_fake,fake)[0]
                d_loss_cum += 0.5 * np.add(d_loss_real, d_loss_fake)

                # Train generator
                g_loss_cum += self.encoder_discriminator.train_on_batch(grids, valid)[0]

            d_loss = d_loss_cum/(step+1)
            g_loss = g_loss_cum/(step+1)

            with train_summary_writer.as_default():
                tf.summary.scalar('loss - ae', loss, step=epoch)
                tf.summary.scalar('accuracy - ae', acc, step=epoch)
                tf.summary.scalar('loss - g', g_loss, step=epoch)
                tf.summary.scalar('loss - d', d_loss, step=epoch)

            # Calculate the accuracies on the validation set
            if val_data is not None:
                loss_val, acc_val, d_loss_val, g_loss_val = \
                    self.validate(val_dataset, val_batch_size)

                with val_summary_writer.as_default():
                    tf.summary.scalar('loss - ae', loss_val, step=epoch)
                    tf.summary.scalar('accuracy - ae', acc_val, step=epoch)
                    tf.summary.scalar('loss - g', g_loss_val, step=epoch)
                    tf.summary.scalar('loss - d', d_loss_val, step=epoch)


    def validate(self, val_dataset, val_batch_size=128):

        # Adversarial ground truths
        valid = np.ones((val_batch_size, 1))
        fake = np.zeros((val_batch_size, 1))

        loss_cum = 0
        acc_cum = 0
        d_loss_cum = 0
        g_loss_cum = 0
        for step, val_grids in enumerate(val_dataset):

            loss, acc = self.autoencoder.evaluate(val_grids, val_grids, verbose=0)
            loss_cum += loss
            acc_cum += acc

            latent_fake = self.encoder.predict(val_grids)
            latent_real = np.random.normal(size=(val_batch_size, self.latent_dim))

            d_loss_real = self.discriminator.evaluate(latent_real, valid, verbose=0)[0]
            d_loss_fake = self.discriminator.evaluate(latent_fake, fake, verbose=0)[0]
            d_loss_cum += 0.5 * np.add(d_loss_real, d_loss_fake)

            g_loss_cum += self.encoder_discriminator.evaluate(val_grids, valid,verbose=0)[0]

        # Average the loss and accuracy over the entire dataset
        loss = loss_cum/(step+1)
        acc = acc_cum/(step+1)
        d_loss = d_loss_cum/(step+1)
        g_loss = g_loss_cum/(step+1)

        return loss, acc, d_loss, g_loss