⌂

Image Randomiser

A Method to Optimize Samples by Downsampling and Langevin Annealing using JS Divergence

This repository contains a Python script that employs Langevin dynamics to process a set of images and save a sample of the optimal images to a specified destination directory.

Python Code: https://github.com/1ssb/Image-Randomiser/blob/main/code.py

Overview

This is an image sampling algorithm that samples a subset of images to optimize sample diversity by minimising the JS divergence between sample distributiuon and source distribution. We employ PyTorch to implement the image processing pipeline and can utilize graphical processing for acceleration if available.

The core problem this code is solving is to select the most diverse images from a dataset using JS divergence loss function after annealing the distributions generated in the samples and the source. It adds noise to the images during the evaluation process using Langevin dynamics schedule to explore different solution space regions and avoid getting trapped in local minima.

Commentry

In this code, we are importing necessary libraries such as os, shutil, numpy, PIL, torch, torchvision, and tqdm. After importing the libraries, we check if a GPU is available and if not, raise an exception. If a GPU is available, we set the device to cuda.

We then define a class called ImageDataset which inherits from the Dataset class in PyTorch. This class takes in a root_dir argument during initialization and stores it as an instance variable. We also create an instance variable called image_paths which is a list of all image file paths in the root_dir that end with .png, .jpg, or .jpeg. The class also has two methods: __len__ and __getitem__. The __len__ method returns the number of image paths in the dataset while the __getitem__ method takes in an index and returns the image at that index after converting it to RGB and transforming it into a tensor.

We then define two functions: select_diverse_images and copy_diverse_images. The first function takes in a source distribution, a value for k, and a batch size. It returns the best subset of images from the source distribution based on their Jensen-Shannon divergence of the random sample and source, iterated over batches. The second function takes in a source path, destination path, value for k, number of evaluations, and batch size. It creates an instance of the ImageDataset class using the source path and creates a DataLoader from it. It then computes the source distribution by flattening each image in the dataset and concatenating them. Next, it selects the best subset of images from the source distribution using the previously defined function and copies those images to the destination path.

Algorithm's overview
Algorithm's overview*

Background

Jensen-Shannon Divergence

The advantage of JS divergence over KL divergence is mostly related to issues with empty probabilities for certain events or bins and how these cause issues with Kullback-Leibler divergence (KL Divergence) and population stability index (PSI). JS divergence uses a mixture probability as a baseline when comparing two distributions. In the discrete versions of PSI and KL divergence, the equations blow up when there are 0 probability events.

JS-divergence is an improved version of KL-divergence with symmetric measures and normalized outcomes. The normalization provides better stability during the loss function optimization.

The equation for JS-divergence between two distributions g and h can be defined as:

JS(g,h) = 0.5 * KL(g:(g+h)/2) + 0.5 * KL(h:(g+h)/2)

where KL represents Kullback-Leibler divergence. From this equation, we can say that JS-divergence computes a symmetrical measure. If we use base-2 logarithm then this measure is a normalized version of KL divergence, with scores between 0 (same) and 1 (completely different).

Langevin Annealing

Langevin annealing is a method that uses Langevin dynamics to optimize a non-convex potential function. It is based on the idea of adding random noise to the motion of particles to simulate the effects of thermal fluctuations. This allows the optimization algorithm to explore different regions of the solution space and avoid getting stuck in local minima. The noise is added at multiple levels using an annealing schedule, which gradually reduces the amount of noise as the optimization progresses.

Usage

At the bottom of the script, we define some variables for the source path, destination path, value for k (number of samples), number of evaluations, and batch size for the divergence calculation. We then call the copy_diverse_images function with these arguments to execute the code. Make sure you use the correct paths and the desired parameters, also you could explore changing the hyperparameter of the Langevin annealing.

Check out the code and email me at Subhransu.Bhattacharjee@anu.edu.au if you encounter any issues.

References

  1. Welling, M., & Teh, Y. W. (2011). Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th International Conference on Machine Learning (ICML-11) (pp. 681-688).
  2. Li, C., Zhu, J., & Zhang, B. (2018). Learning energy-based models with exponential family Langevin dynamics. arXiv preprint arXiv:1811.12359.
  3. Pascanu, R., & Bengio, Y. (2014). Revisiting natural gradient for deep networks. arXiv preprint arXiv:1301.3584.
  4. Lin, J. (1991). Divergence measures based on the Shannon entropy. IEEE Transactions on Information Theory, 37(1), 145-151.
  5. Weng, L. (2019). From GAN to WGAN. arXiv preprint arXiv:1904.08994.
*Langevin annealing images is from Yang Song's blog.

Cite this page

To cite this work, you can use the following format:

Bhattacharjee, S. S. (2023). Image Randomiser: A Method to Optimize Samples by Downsampling and Langevin Annealing using JS Divergence. Retrieved from https://github.com/1ssb/Image-Randomiser

© 2023 Subhransu S. Bhattacharjee