Example_Notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Unsupervised Lensing\n",
    "A PyTorch-based tool for Unsupervised Deep Learning applications in strong lensing cosmology\n",
    "\n",
    "Currently supported models:\n",
    "\n",
    "* Adversarial Autoencoder\n",
    "\n",
    "* Convolutional Variational Autoencoder\n",
    "\n",
    "* Deep Convolutional Autoencoder\n",
    "\n",
    "* Restricted Boltzmann Machine\n",
    "\n",
    "### Example 1: Data pre-processing\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from unsupervised_lensing.utils.data_preprocess import process\n",
    "\n",
    "data = process(data_path='./Data',    # Path to your input data images\n",
    "               batch_size = 100)      # Batch Size\n",
    "\n",
    "np.save('no_sub_train.npy', data)     # Save the training data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ">**[NOTE]**\n",
    ">torch data loader expects the input images to be in subfolders, so arrange your input directory accordingly\n",
    "\n",
    ">Example: \"./Data/no_sub/no_sub_1.png\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 2: Adversarial Autoencoder\n",
    "\n",
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from unsupervised_lensing.models import Adversarial_AE\n",
    "from unsupervised_lensing.models.AAE_Nets import *\n",
    "from unsupervised_lensing.utils import loss_plotter as plt\n",
    "from unsupervised_lensing.utils.EMD_Lensing import EMD\n",
    "\n",
    "# Model Training\n",
    "out = Adversarial_AE.train(data_path='./Data/no_sub_train.npy', # Path to training dataset\n",
    "                           epochs=100,\n",
    "                           learning_rate=2e-3,\n",
    "                           optimizer='Adam',\n",
    "                           checkpoint_path='./Weights',         # Path to store model weights\n",
    "                           pretrain=True,                       # Set True for transfer learning\n",
    "                           pretrain_mode='transfer',            # change to 'continue' to load weights from 'checkpoint_path' directory\n",
    "                           pretrain_model='A')                  # Select the model for transfer learning     \n",
    "\n",
    "# Plot the training loss\n",
    "plt.plot_loss(out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model Validation\n",
    "recon_loss = Adversarial_AE.evaluate(data_path='./Data/no_sub_test.npy', # Path to validation dataset\n",
    "                                     checkpoint_path='./Weights',        # Path to model weights\n",
    "                                     out_path='./Results')               # Path to store reconstructed samples\n",
    "\n",
    "# Plot the reconstruction loss\n",
    "plt.plot_dist(recon_loss)\n",
    "\n",
    "# Calculate Wasserstein distance\n",
    "print(EMD(data_path='./Data/no_sub_test.npy', recon_path='./Results/Recon_samples.npy'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 3: Variational Autoencoder\n",
    "\n",
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from unsupervised_lensing.models import Variational_AE\n",
    "from unsupervised_lensing.models.VAE_Nets import *\n",
    "\n",
    "# Model Training\n",
    "out = Variational_AE.train(data_path='./Data/no_sub_train.npy', # Path to training dataset\n",
    "                           epochs=100,\n",
    "                           learning_rate=2e-3,\n",
    "                           optimizer='Adam',\n",
    "                           checkpoint_path='./Weights',         # Path to store model weights\n",
    "                           pretrain=True,                       # Set True for transfer learning\n",
    "                           pretrain_mode='transfer',            # change to 'continue' to load weights from 'checkpoint_path' directory\n",
    "                           pretrain_model='A')                  # Select the model for transfer learning     \n",
    "\n",
    "# Plot the training loss\n",
    "plt.plot_loss(out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model Validation\n",
    "recon_loss = Variational_AE.evaluate(data_path='./Data/no_sub_test.npy', # Path to validation dataset\n",
    "                                     checkpoint_path='./Weights',        # Path to model weights\n",
    "                                     out_path='./Results')               # Path to store reconstructed samples\n",
    "\n",
    "# Plot the reconstruction loss\n",
    "plt.plot_dist(recon_loss)\n",
    "\n",
    "# Calculate Wasserstein distance\n",
    "print(EMD(data_path='./Data/no_sub_test.npy', recon_path='./Results/Recon_samples.npy'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 4: Deep Convolutional Autoencoder\n",
    "\n",
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from unsupervised_lensing.models import Convolutional_AE\n",
    "from unsupervised_lensing.models.DCAE_Nets import *\n",
    "\n",
    "# Model Training\n",
    "out = Convolutional_AE.train(data_path='./Data/no_sub_train.npy', # Path to training dataset\n",
    "                             epochs=100,\n",
    "                             learning_rate=2e-3,\n",
    "                             optimizer='Adam',\n",
    "                             checkpoint_path='./Weights',         # Path to store model weights\n",
    "                             pretrain=True,                       # Set True for transfer learning\n",
    "                             pretrain_mode='transfer',            # change to 'continue' to load weights from 'checkpoint_path' directory\n",
    "                             pretrain_model='A')                  # Select the model for transfer learning     \n",
    "\n",
    "# Plot the training loss\n",
    "plt.plot_loss(out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model Validation\n",
    "recon_loss = Convolutional_AE.evaluate(data_path='./Data/no_sub_test.npy', # Path to validation dataset\n",
    "                                       checkpoint_path='./Weights',        # Path to model weights\n",
    "                                       out_path='./Results')               # Path to store reconstructed samples\n",
    "\n",
    "# Plot the reconstruction loss\n",
    "plt.plot_dist(recon_loss)\n",
    "\n",
    "# Calculate Wasserstein distance\n",
    "print(EMD(data_path='./Data/no_sub_test.npy', recon_path='./Results/Recon_samples.npy'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 5: Restricted Boltzmann Machine\n",
    "\n",
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from unsupervised_lensing.models import RBM_Model\n",
    "from unsupervised_lensing.models.RBM_Nets import *\n",
    "\n",
    "# Model Training\n",
    "out = RBM_Model.train(data_path='./Data/no_sub_train.npy', # Path to training dataset\n",
    "                      epochs=100,\n",
    "                      learning_rate=2e-3,\n",
    "                      optimizer='Adam',\n",
    "                      checkpoint_path='./Weights',         # Path to store model weights\n",
    "                      pretrain=True,                       # Set True for transfer learning\n",
    "                      pretrain_mode='transfer',            # change to 'continue' to load weights from 'checkpoint_path' directory\n",
    "                      pretrain_model='A')                  # Select the model for transfer learning     \n",
    "\n",
    "# Plot the training loss\n",
    "plt.plot_loss(out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model Validation\n",
    "recon_loss = RBM_Model.evaluate(data_path='./Data/no_sub_test.npy', # Path to validation dataset\n",
    "                                checkpoint_path='./Weights',        # Path to model weights\n",
    "                                out_path='./Results')               # Path to store reconstructed samples\n",
    "\n",
    "# Plot the reconstruction loss\n",
    "plt.plot_dist(recon_loss)\n",
    "\n",
    "# Calculate Wasserstein distance\n",
    "print(EMD(data_path='./Data/no_sub_test.npy', recon_path='./Results/Recon_samples.npy'))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}