Variational autoencoder for anomaly detection

PyPI - Python Version GitHub Workflow Status

Pytorch/TF1 implementation of Variational AutoEncoder for anomaly detection following the paper Variational Autoencoder based Anomaly Detection using Reconstruction Probability by Jinwon An, Sungzoon Cho

How to install

pip install vae-anomaly-detection

How To Train a Model

  • Define your dataset into dataset.py and put in output into the function get_dataset
  • Eventually change encoder and decoder inside VAE.py to fits your data layout
  • Run in a terminal python train.py and specify required at least --input-size (pass -h to see all optional parameters)
  • Trained model, parameters and Tensorboard log goes into the folder run/{id} where {id} is an integer from 0 to +inf
  • After the model training run tensorboard --logdir=run to check all the training results

How to make predictions:

Once the model is trained (suppose for simplicity that it is under run/0/ ) just load and predict with this code snippet:

import torch

#load X_test
model = VAEAnomaly(input_size=50, latent_size=32)
# could load input_size and latent_size also 
# from run/0/train_config.yaml
model.load_state_dict(torch.load('run/0/model.pth'))
# load saved parameters from a run
outliers = model.is_anomaly(X_test)