anomaly_detection_using_vae.py

# -*- coding: utf-8 -*-
"""Anomaly Detection using VAE

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/17eiv6aywFmNnXjrPhgmov5QT8TLdOg1r

# Anomaly Detection
Anomaly detection is an unsupervised pattern recognition task that can be defined under different statistical models.
Given a set of training samples containing no anomalies, the goal of anomaly detection is to design or learn a feature representation, that captures “normal” appearance patterns.

***Here we are using a generative models technique called Variational Autoencoders (VAE) to do Anomaly Detection.***

# **variational autoencoder (VAE)**
A variational autoencoder (VAE) provides a probabilistic manner for describing an observation in latent space. Thus, rather than building an encoder which outputs a single value to describe each latent state attribute, we'll formulate our encoder to describe a probability distribution for each latent attribute.
"""

from keras.layers import Lambda, Input, Dense
from keras.models import Model
from keras.datasets import mnist
from keras.losses import mse, binary_crossentropy
from keras.utils import plot_model
from keras import backend as K

import numpy as np
import matplotlib.pyplot as plt
import argparse
import os

def sampling(args):
    z_mean, z_log_var = args
    batch = K.shape(z_mean)[0]
    dim = K.int_shape(z_mean)[1]    
    epsilon = K.random_normal(shape=(batch, dim))  # by default, random_normal has mean=0 and std=1.0
    return z_mean + K.exp(0.5 * z_log_var) * epsilon

# MNIST dataset
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# get data for one digit "1"

print(x_train.shape)
print(y_train.shape)

dig_class = 1

indexes = []
for i, j in enumerate(y_train):
    if j == dig_class:
        indexes.append(i)
        

x_train_t = x_train[indexes]

indexes = []
for i, j in enumerate(y_test):
    if j == dig_class:
        indexes.append(i)


x_test_t = x_test[indexes]


x_train = x_train_t
x_test = x_test_t
print(x_train.shape)
print(x_test.shape)

# reshape and normalization
image_size = x_train.shape[1]
original_dim = image_size * image_size
x_train = np.reshape(x_train, [-1, original_dim])
x_test = np.reshape(x_test, [-1, original_dim])
x_train = x_train.astype('float32') / 255
x_test = x_test.astype('float32') / 255

# network parameters and learning parameters
input_shape = (original_dim, )
intermediate_dim = 512
batch_size = 128
latent_dim = 2
epochs = 50

# encoder model
inputs = Input(shape=input_shape, name='encoder_input')
x = Dense(intermediate_dim, activation='relu')(inputs)
z_mean = Dense(latent_dim, name='z_mean')(x)
z_log_var = Dense(latent_dim, name='z_log_var')(x)

# sampling 
z = Lambda(sampling, output_shape=(latent_dim,), name='z')([z_mean, z_log_var]) 

# instantiate encoder model
encoder = Model(inputs, [z_mean, z_log_var, z], name='encoder')
encoder.summary()

# plot_model(encoder, to_file='vae_mlp_encoder.png', show_shapes=True)

# decoder model
latent_inputs = Input(shape=(latent_dim,), name='z_sampling')
x = Dense(intermediate_dim, activation='relu')(latent_inputs)
outputs = Dense(original_dim, activation='sigmoid')(x)

# instantiate decoder model
decoder = Model(latent_inputs, outputs, name='decoder')
decoder.summary()

# plot_model(decoder, to_file='vae_mlp_decoder.png', show_shapes=True)

# instantiate VAE model
outputs = decoder(encoder(inputs)[2])
vae = Model(inputs, outputs, name='vae')
vae.summary()

# VAE loss
# reconstruction_loss = mse(inputs, outputs)
reconstruction_loss = binary_crossentropy(inputs, outputs)
reconstruction_loss *= original_dim
kl_loss = 1 + z_log_var - K.square(z_mean) - K.exp(z_log_var)
kl_loss = K.sum(kl_loss, axis=-1)
kl_loss *= -0.5
vae_loss = K.mean(reconstruction_loss + kl_loss)
# vae_loss = K.mean(reconstruction_loss)

# vae.compile(optimizer='adam', loss=vae_loss)
vae.add_loss(vae_loss)
vae.compile(optimizer='adam')
vae.summary()

# Learning
# epochs = 50
history = vae.fit(x_train, epochs=epochs, batch_size=batch_size, validation_data=(x_test, None))

vae.save_weights('vae_mnist.h5')

# plot loss history
loss = history.history['loss']
val_loss = history.history['val_loss']
epochs = range(len(loss))

plt.figure()
plt.plot(epochs, loss, 'bo', label='Training loss')
plt.plot(epochs, val_loss, 'b', label='Validation loss')
plt.title('Training and validation loss')
plt.legend()
plt.show()

# Visualization of latent space
z_mean, _, _ = encoder.predict(x_test, batch_size=batch_size)
plt.figure(figsize=(12, 10))
plt.scatter(z_mean[:, 0], z_mean[:, 1])
plt.xlabel("z[0]")
plt.ylabel("z[1]")
plt.title('Test Data Latent Space')
plt.show()

# Visualization of latent space
z_mean, _, _ = encoder.predict(x_train, batch_size=batch_size)
plt.figure(figsize=(12, 10))
plt.scatter(z_mean[:, 0], z_mean[:, 1])
plt.xlabel("z[0]")
plt.ylabel("z[1]")
plt.title('Train Data Latent Space')
plt.show()

# display a 30x30 2D manifold of digits
    n = 15
    digit_size = 28
    figure = np.zeros((digit_size * n, digit_size * n))
    grid_x = np.linspace(-3, 4, n)
    grid_y = np.linspace(-4, 4, n)[::-1]

    for i, yi in enumerate(grid_y):
        for j, xi in enumerate(grid_x):
            z_sample = np.array([[xi, yi]])
            x_decoded = decoder.predict(z_sample)
            digit = x_decoded[0].reshape(digit_size, digit_size)
            figure[i * digit_size: (i + 1) * digit_size,
                   j * digit_size: (j + 1) * digit_size] = digit

    plt.figure(figsize=(10, 10))
    start_range = digit_size // 2
    end_range = n * digit_size + start_range + 1
    pixel_range = np.arange(start_range, end_range, digit_size)
    sample_range_x = np.round(grid_x, 1)
    sample_range_y = np.round(grid_y, 1)
    plt.xticks(pixel_range, sample_range_x)
    plt.yticks(pixel_range, sample_range_y)
    plt.xlabel("z[0]")
    plt.ylabel("z[1]")
    plt.imshow(figure, cmap='gray')
    plt.show()

# MNIST dataset
(x_train, y_train), (x_test, y_test) = mnist.load_data()

image_size = x_train.shape[1]
original_dim = image_size * image_size
x_train = np.reshape(x_train, [-1, original_dim])
x_test = np.reshape(x_test, [-1, original_dim])
x_train = x_train.astype('float32') / 255
x_test = x_test.astype('float32') / 255

"""# Visualization of Latent Space
When we supply the whole data to the trained VAE, we can see that result is separable
the anomaly digits (digits which is not "one") are outside the distribution of normal latent space.
"""

# Visualization of latent space
z_mean,z_log_var, _ = encoder.predict(x_test, batch_size=batch_size)
plt.figure(figsize=(12, 10))
plt.scatter(z_mean[:, 0], z_mean[:, 1], c=y_test)
plt.colorbar()
plt.set_cmap('tab10')
plt.xlabel("z[0]")
plt.ylabel("z[1]")
plt.title('Latent Space for All Data')
plt.show()

# Visualization of latent space
plt.figure(figsize=(12, 10))
plt.scatter(z_log_var[:, 0], z_log_var[:, 1], c=y_test)
plt.colorbar()
plt.set_cmap('tab10')
plt.xlabel("z_log_var[0]")
plt.ylabel("z_log_var[1]")
plt.title('Latent Space for All Data')
plt.show()

# test reconstruction for one digit

i = 14

digit_size = 28


digit = x_train[i].reshape(digit_size, digit_size)

figure = digit
plt.figure(figsize=(5, 5))
plt.imshow(figure, cmap='gray')
plt.show()



z_sample = np.array(z_mean)
x_decoded = decoder.predict(z_sample)
digit = x_decoded[i].reshape(digit_size, digit_size)


figure = digit
plt.figure(figsize=(5, 5))
plt.imshow(figure, cmap='gray')
plt.show()

# test reconstruction
    n = 20
    digit_size = 28
    figure = np.zeros((digit_size * 2, digit_size * n))
    grid_x = np.linspace(0, n, n)


    for j, xi in enumerate(grid_x):
        z_sample = np.array(z_mean)
        x_decoded = decoder.predict(z_sample)
        digit = x_train[j].reshape(digit_size, digit_size)
        figure[0 * digit_size: (0 + 1) * digit_size,
               j * digit_size: (j + 1) * digit_size] = digit
        digit = x_decoded[j].reshape(digit_size, digit_size)
        figure[1 * digit_size: (1 + 1) * digit_size,
               j * digit_size: (j + 1) * digit_size] = digit

    plt.figure(figsize=(n, n))
    
    plt.imshow(figure, cmap='gray')
    plt.show()