Task 2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Q. Using a deep learning algorithm of your choice, learn the representation of dark matter in the strong gravitational lensing images provided using PyTorch.\n",
    "The goal is to use this learned representation for anomaly detection, hence you should pick the most appropriate algorithm and discuss your strategy.   \n",
    "   \n",
    "A. Considering the purpose is anomaly detection, I choose Autoencoder as the most appropriate method, as discussed in [1]. I think Conditional VAE (CVAE)[5] especially suits this purpose.   \n",
    "Deep neural network architecture is known to function relatively better in tasks about images than MLP(multilayer perception), and in the early stage Convolutional AutoEncoder(hereafter CAE)[2] and Variational AutoEncoder(VAE)[8] was the most competent method.\n",
    "Of cource, it is said that CNN architecture in general performs well, speaking of image processing.  \n",
    "CAE incorpolates CNN's algorithm into itself, and according to [3],[4] (written in Japnanese), it clearly shows that CAE overwhelms other MLP algorithms.   \n",
    "VAE incorpolates probability distribution into latent variable z, and it also performs as well as CAE does.\n",
    "These days, lots of methods derived from these, and I noticed CVAE is the best, which add supervising features to VAE. it is compared with lots of other autoencoder methods on purpose of analysis on astronomical images[6], and prove that the latent variables obtained with CVAE preserve more information and are more discriminative towards real astronomical transients.(SCAE[7] did well as well, but slightly CVAE has an advantage of it)    \n",
    "This paper was published in April 2018, and it is the most recent paper applying autoencoder to processing astronomical images.(According to my survey on Google Scholor)\n",
    "In addition to that, CVAE is implementable on laptop, whereas other methods sometimes need a bulk of computer properties(such as [9].)\n",
    "\n",
    "Using such autoencoders, anomaly detection will be easy.(in future work I can impliment heat map which clearly tells anomalies from norms [10] )\n",
    "\n",
    "According to Hold-out method, I devided images into train(3001 images) and test(2000 images).  \n",
    "Directory structures are below. you should put lenses directory at the same place as this notebook. \n",
    "- lenses/train/sub/(3001 images)\n",
    "- lenses/train/no_sub/(3001 images)\n",
    "- lenses/test/sub/(2000 images)\n",
    "- lenses/test/no_sub/(2000 images)\n",
    "\n",
    "About optimizer, I use Adam because this is most popular optimizers' father and since this is not that complicated, the convergence speed is fast.\n",
    "\n",
    "in inplementation, I refered to this site:  \n",
    "https://graviraja.github.io/conditionalvae/#\n",
    "\n",
    "\n",
    "----\n",
    "\n",
    "\n",
    "Of cource, CNN could be used because this problem can be taken as a classification problem(so no_sub images provided)\n",
    "~~In case, I also implement CNN, which was written at the end of this notebook.(I used ResNet, as in [11])~~\n",
    "I couldn't implement CNN case in time.\n",
    "\n",
    "\n",
    "\n",
    "----\n",
    "\n",
    "[1]Deep Learning the Morphology of Dark Matter Substructure, pp8  \n",
    "https://arxiv.org/pdf/1909.07346.pdf\n",
    "\n",
    "[2]Deep Clustering with Convolutional Autoencoders  \n",
    "https://link.springer.com/chapter/10.1007/978-3-319-70096-0_39\n",
    "\n",
    "[3]PytorchによるAutoEncoder Familyの実装(Implementation of various Autoencoders using PyTorch)  \n",
    "https://elix-tech.github.io/ja/2016/07/17/autoencoder.html\n",
    "\n",
    "[4]様々なオートエンコーダによる異常検知(Anomaly Detection with lots of Autoencoders)  \n",
    "https://sinyblog.com/deaplearning/auto_encoder_001/\n",
    "\n",
    "[5]Semi-Supervised Learning with Deep Generative Models (CVAE)  \n",
    "https://arxiv.org/abs/1406.5298\n",
    "\n",
    "[6]Latent representations of transient candidates from an astronomical image difference pipeline using Variational Autoencoders  \n",
    "https://pdfs.semanticscholar.org/57ce/a9520008706b8cb190d94513e695068210cd.pdf\n",
    "\n",
    "[7]Convolutional Sparse Autoencoders for Image Classification (SCAE)  \n",
    "https://ieeexplore.ieee.org/document/7962256\n",
    "\n",
    "[8]Auto-Encoding Variational Bayes(VAE)  \n",
    "https://arxiv.org/pdf/1312.6114.pdf\n",
    "\n",
    "[9]Improving Variational Autoencoder with Deep Feature Consistent and Generative Adversarial Training (GAN)  \n",
    "https://arxiv.org/pdf/1906.01984.pdf\n",
    "\n",
    "[10]Anomaly Manufacturing Product Detection using Unregularized Anomaly Score on Deep Generative Models  \n",
    "https://confit.atlas.jp/guide/event-img/jsai2018/2A1-03/public/pdf?type=in  \n",
    "https://qiita.com/shinmura0/items/811d01384e20bfd1e035\n",
    "\n",
    "[11] Deep Learning the Morphology of Dark Matter Substructure https://arxiv.org/abs/1909.07346"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1: CVAE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torchvision import datasets, transforms, utils\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "from skimage import io, transform\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 1         # number of data points in each batch\n",
    "N_EPOCHS = 1           # times to run the model on complete data\n",
    "INPUT_DIM = 150 * 150     # size of each input\n",
    "HIDDEN_DIM = 4        # hidden dimension\n",
    "LATENT_DIM = 15         # latent vector dimension\n",
    "N_CLASSES = 2          # number of classes in the data\n",
    "lr = 1e-3               # learning rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_transform = transforms.Compose([transforms.RandomSizedCrop(INPUT_DIM), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])\n",
    "\n",
    "train_dataset = datasets.ImageFolder(root='./lenses/train', #3001 images each\n",
    "                                           transform=data_transform)\n",
    "test_dataset = datasets.ImageFolder(root='./lenses/test', #2000 images each\n",
    "                                           transform=data_transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_iterator = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
    "test_iterator = DataLoader(test_dataset, batch_size=BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "def idx2onehot(idx, n=N_CLASSES):\n",
    "\n",
    "    assert idx.shape[1] == 1\n",
    "    assert torch.max(idx).item() < n\n",
    "\n",
    "    onehot = torch.zeros(idx.size(0), n)\n",
    "    onehot.scatter_(1, idx.data, 1)\n",
    "\n",
    "    return onehot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    ''' This the encoder part of VAE\n",
    "    '''\n",
    "    def __init__(self, input_dim, hidden_dim, latent_dim, n_classes):\n",
    "        '''\n",
    "        Args:\n",
    "            input_dim: A integer indicating the size of input \n",
    "            hidden_dim: A integer indicating the size of hidden dimension.\n",
    "            latent_dim: A integer indicating the latent size.\n",
    "            n_classes: A integer indicating the number of classes. (dimension of one-hot representation of labels)\n",
    "        '''\n",
    "        super().__init__()\n",
    "\n",
    "        self.linear = nn.Linear(input_dim + n_classes, hidden_dim)\n",
    "        self.mu = nn.Linear(hidden_dim, latent_dim)\n",
    "        self.var = nn.Linear(hidden_dim, latent_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # x is of shape [batch_size, input_dim + n_classes]\n",
    "\n",
    "        hidden = F.relu(self.linear(x))\n",
    "        # hidden is of shape [batch_size, hidden_dim]\n",
    "\n",
    "        # latent parameters\n",
    "        mean = self.mu(hidden)\n",
    "        # mean is of shape [batch_size, latent_dim]\n",
    "        log_var = self.var(hidden)\n",
    "        # log_var is of shape [batch_size, latent_dim]\n",
    "\n",
    "        return mean, log_var\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    ''' This the decoder part of VAE\n",
    "    '''\n",
    "    def __init__(self, latent_dim, hidden_dim, output_dim, n_classes):\n",
    "        '''\n",
    "        Args:\n",
    "            latent_dim: A integer indicating the latent size.\n",
    "            hidden_dim: A integer indicating the size of hidden dimension.\n",
    "            output_dim: A integer indicating the size of output \n",
    "            n_classes: A integer indicating the number of classes. (dimension of one-hot representation of labels)\n",
    "        '''\n",
    "        super().__init__()\n",
    "\n",
    "        self.latent_to_hidden = nn.Linear(latent_dim + n_classes, hidden_dim)\n",
    "        self.hidden_to_out = nn.Linear(hidden_dim, output_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # x is of shape [batch_size, latent_dim + num_classes]\n",
    "        x = F.relu(self.latent_to_hidden(x))\n",
    "        # x is of shape [batch_size, hidden_dim]\n",
    "        generated_x = F.sigmoid(self.hidden_to_out(x))\n",
    "        # x is of shape [batch_size, output_dim]\n",
    "\n",
    "        return generated_x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CVAE(nn.Module):\n",
    "    ''' This the VAE, which takes a encoder and decoder.\n",
    "    '''\n",
    "    def __init__(self, input_dim, hidden_dim, latent_dim, n_classes):\n",
    "        '''\n",
    "        Args:\n",
    "            input_dim: A integer indicating the size of input.\n",
    "            hidden_dim: A integer indicating the size of hidden dimension.\n",
    "            latent_dim: A integer indicating the latent size.\n",
    "            n_classes: A integer indicating the number of classes. (dimension of one-hot representation of labels)\n",
    "        '''\n",
    "        super().__init__()\n",
    "\n",
    "        self.encoder = Encoder(input_dim, hidden_dim, latent_dim, n_classes)\n",
    "        self.decoder = Decoder(latent_dim, hidden_dim, input_dim, n_classes)\n",
    "\n",
    "    def forward(self, x, y):\n",
    "\n",
    "        x = torch.cat((x, y), dim=1)\n",
    "\n",
    "        # encode\n",
    "        z_mu, z_var = self.encoder(x)\n",
    "\n",
    "        # sample from the distribution having latent parameters z_mu, z_var\n",
    "        # reparameterize\n",
    "        std = torch.exp(z_var / 2)\n",
    "        eps = torch.randn_like(std)\n",
    "        x_sample = eps.mul(std).add_(z_mu)\n",
    "\n",
    "        z = torch.cat((x_sample, y), dim=1)\n",
    "\n",
    "        # decode\n",
    "        generated_x = self.decoder(z)\n",
    "\n",
    "        return generated_x, z_mu, z_var"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_loss(x, reconstructed_x, mean, log_var):\n",
    "    # reconstruction loss\n",
    "    RCL = F.binary_cross_entropy(reconstructed_x, x, size_average=False)\n",
    "    # kl divergence loss\n",
    "    KLD = -0.5 * torch.sum(1 + log_var - mean.pow(2) - log_var.exp())\n",
    "\n",
    "    return RCL + KLD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = CVAE(INPUT_DIM, HIDDEN_DIM, LATENT_DIM, N_CLASSES)\n",
    "optimizer = optim.Adam(model.parameters(), lr=lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train():\n",
    "    # set the train mode\n",
    "    model.train()\n",
    "\n",
    "    # loss of the epoch\n",
    "    train_loss = 0\n",
    "\n",
    "    for i, (x, y) in enumerate(train_iterator):\n",
    "        # reshape the data into \n",
    "        x = x.view(-1,  INPUT_DIM*3)\n",
    "        x = x.to(device)\n",
    "\n",
    "        # convert y into one-hot encoding\n",
    "        y = idx2onehot(y.view(-1, 1))\n",
    "        y = y.to(device)\n",
    "\n",
    "        # update the gradients to zero\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # forward pass\n",
    "        reconstructed_x, z_mu, z_var = model(x, y)\n",
    "\n",
    "        # loss\n",
    "        loss = calculate_loss(x, reconstructed_x, z_mu, z_var)\n",
    "\n",
    "        # backward pass\n",
    "        loss.backward()\n",
    "        train_loss += loss.item()\n",
    "\n",
    "        # update the weights\n",
    "        optimizer.step()\n",
    "\n",
    "    return train_loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test():\n",
    "    # set the evaluation mode\n",
    "    model.eval()\n",
    "\n",
    "    # test loss for the data\n",
    "    test_loss = 0\n",
    "\n",
    "    # we don't need to track the gradients, since we are not updating the parameters during evaluation / testing\n",
    "    with torch.no_grad():\n",
    "        for i, (x, y) in enumerate(test_iterator):\n",
    "            # reshape the data\n",
    "            x = x.view(-1, INPUT_DIM*3) #TASK: modifying shapes of the dataset, data shapes in one batch size must be the same https://nodaki.hatenablog.com/entry/2018/07/18/001851\n",
    "            x = x.to(device)\n",
    "\n",
    "            # convert y into one-hot encoding\n",
    "            y = idx2onehot(y.view(-1, 1))\n",
    "            y = y.to(device)\n",
    "\n",
    "            # forward pass\n",
    "            reconstructed_x, z_mu, z_var = model(x, y)\n",
    "\n",
    "            # loss\n",
    "            loss = calculate_loss(x, reconstructed_x, z_mu, z_var)\n",
    "            test_loss += loss.item()\n",
    "\n",
    "    return test_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "invalid argument 0: Sizes of tensors must match except in dimension 1. Got 22500 and 1 in dimension 0 at /Users/distiller/project/conda/conda-bld/pytorch_1579022061893/work/aten/src/TH/generic/THTensor.cpp:612",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-54-4227fd0f2dc4>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0me\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mN_EPOCHS\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m     \u001b[0mtrain_loss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      6\u001b[0m     \u001b[0mtest_loss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtest\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-52-a71cedf56391>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m()\u001b[0m\n\u001b[1;32m     19\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     20\u001b[0m         \u001b[0;31m# forward pass\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 21\u001b[0;31m         \u001b[0mreconstructed_x\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mz_mu\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mz_var\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     22\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     23\u001b[0m         \u001b[0;31m# loss\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/anaconda3/envs/autolens/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    530\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    531\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 532\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    533\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    534\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-49-20b2a6961901>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x, y)\u001b[0m\n\u001b[1;32m     17\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 19\u001b[0;31m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     20\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m         \u001b[0;31m# encode\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: invalid argument 0: Sizes of tensors must match except in dimension 1. Got 22500 and 1 in dimension 0 at /Users/distiller/project/conda/conda-bld/pytorch_1579022061893/work/aten/src/TH/generic/THTensor.cpp:612"
     ]
    }
   ],
   "source": [
    "best_test_loss = float('inf')\n",
    "\n",
    "for e in range(N_EPOCHS):\n",
    "\n",
    "    train_loss = train()\n",
    "    test_loss = test()\n",
    "\n",
    "    train_loss /= len(train_dataset)\n",
    "    test_loss /= len(test_dataset)\n",
    "\n",
    "    print(f'Epoch {e}, Train Loss: {train_loss:.2f}, Test Loss: {test_loss:.2f}')\n",
    "\n",
    "    if best_test_loss > test_loss:\n",
    "        best_test_loss = test_loss\n",
    "        patience_counter = 1\n",
    "    else:\n",
    "        patience_counter += 1\n",
    "\n",
    "    if patience_counter > 3:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a random latent vector\n",
    "z = torch.randn(1, LATENT_DIM).to(device)\n",
    "\n",
    "# pick randomly 1 class, for which we want to generate the data\n",
    "y = torch.randint(0, N_CLASSES, (1, 1)).to(dtype=torch.long)\n",
    "print(f'Generating a {y.item()}')\n",
    "\n",
    "y = idx2onehot(y).to(device, dtype=z.dtype)\n",
    "z = torch.cat((z, y), dim=1)\n",
    "\n",
    "reconstructed_img = model.decoder(z)\n",
    "img = reconstructed_img.view(150, 150).data\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(img, cmap='gray')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### below are implemented otherwise and nothing to do with CVAE.  \n",
    "\n",
    "\n",
    "# 2: VAE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torchvision import datasets, transforms, utils\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "from skimage import io, transform\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 1        # number of data points in each batch\n",
    "N_EPOCHS = 2           # times to run the model on complete data\n",
    "INPUT_DIM = 150 * 150     # size of each input\n",
    "HIDDEN_DIM = 4        # hidden dimension\n",
    "LATENT_DIM = 15         # latent vector dimension\n",
    "lr = 1e-3               # learning rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_transform = transforms.Compose([transforms.RandomSizedCrop(INPUT_DIM), transforms.ToTensor()])\n",
    "\n",
    "train_dataset = datasets.ImageFolder(root='./lenses/train', #3001 images each\n",
    "                                           transform=data_transform)\n",
    "test_dataset = datasets.ImageFolder(root='./lenses/test', #2000 images each\n",
    "                                           transform=data_transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_iterator = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
    "test_iterator = DataLoader(test_dataset, batch_size=BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    ''' This the encoder part of VAE\n",
    "    '''\n",
    "    def __init__(self, input_dim, hidden_dim, z_dim):\n",
    "        '''\n",
    "        Args:\n",
    "            input_dim: A integer indicating the size of input (in case of MNIST 28 * 28).\n",
    "            hidden_dim: A integer indicating the size of hidden dimension.\n",
    "            z_dim: A integer indicating the latent dimension.\n",
    "        '''\n",
    "        super().__init__()\n",
    "\n",
    "        self.linear = nn.Linear(input_dim, hidden_dim)\n",
    "        self.mu = nn.Linear(hidden_dim, z_dim)\n",
    "        self.var = nn.Linear(hidden_dim, z_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # x is of shape [batch_size, input_dim]\n",
    "\n",
    "        hidden = F.relu(self.linear(x))\n",
    "        # hidden is of shape [batch_size, hidden_dim]\n",
    "        z_mu = self.mu(hidden)\n",
    "        # z_mu is of shape [batch_size, latent_dim]\n",
    "        z_var = self.var(hidden)\n",
    "        # z_var is of shape [batch_size, latent_dim]\n",
    "\n",
    "        return z_mu, z_var"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    ''' This the decoder part of VAE\n",
    "    '''\n",
    "    def __init__(self, z_dim, hidden_dim, output_dim):\n",
    "        '''\n",
    "        Args:\n",
    "            z_dim: A integer indicating the latent size.\n",
    "            hidden_dim: A integer indicating the size of hidden dimension.\n",
    "            output_dim: A integer indicating the output dimension (in case of MNIST it is 28 * 28)\n",
    "        '''\n",
    "        super().__init__()\n",
    "\n",
    "        self.linear = nn.Linear(z_dim, hidden_dim)\n",
    "        self.out = nn.Linear(hidden_dim, output_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # x is of shape [batch_size, latent_dim]\n",
    "\n",
    "        hidden = F.relu(self.linear(x))\n",
    "        # hidden is of shape [batch_size, hidden_dim]\n",
    "\n",
    "        predicted = torch.sigmoid(self.out(hidden))\n",
    "        # predicted is of shape [batch_size, output_dim]\n",
    "        return predicted\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VAE(nn.Module):\n",
    "    def __init__(self, enc, dec):\n",
    "        ''' This the VAE, which takes a encoder and decoder.\n",
    "        '''\n",
    "        super().__init__()\n",
    "\n",
    "        self.enc = enc\n",
    "        self.dec = dec\n",
    "\n",
    "    def forward(self, x):\n",
    "        # encode\n",
    "        z_mu, z_var = self.enc(x)\n",
    "\n",
    "        # sample from the distribution having latent parameters z_mu, z_var\n",
    "        # reparameterize\n",
    "        std = torch.exp(z_var / 2)\n",
    "        eps = torch.randn_like(std)\n",
    "        x_sample = eps.mul(std).add_(z_mu)\n",
    "\n",
    "        # decode\n",
    "        predicted = self.dec(x_sample)\n",
    "        return predicted, z_mu, z_var\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "# encoder\n",
    "encoder = Encoder(INPUT_DIM, HIDDEN_DIM, LATENT_DIM)\n",
    "\n",
    "# decoder\n",
    "decoder = Decoder(LATENT_DIM, HIDDEN_DIM, INPUT_DIM)\n",
    "\n",
    "# vae\n",
    "model = VAE(encoder, decoder).to(device)\n",
    "\n",
    "# optimizer\n",
    "optimizer = optim.Adam(model.parameters(), lr=lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train():\n",
    "    # set the train mode\n",
    "    model.train()\n",
    "\n",
    "    # loss of the epoch\n",
    "    train_loss = 0\n",
    "\n",
    "    for i, (x, _) in enumerate(train_iterator):\n",
    "        # reshape the data into [batch_size, INPUT_DIM]\n",
    "        x = x.view(-1, INPUT_DIM)\n",
    "        x = x.to(device)\n",
    "\n",
    "        # update the gradients to zero\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # forward pass\n",
    "        x_sample, z_mu, z_var = model(x)\n",
    "\n",
    "        # reconstruction loss\n",
    "        recon_loss = F.binary_cross_entropy(x_sample, x, size_average=False)\n",
    "\n",
    "        # kl divergence loss\n",
    "        kl_loss = 0.5 * torch.sum(torch.exp(z_var) + z_mu**2 - 1.0 - z_var)\n",
    "\n",
    "        # total loss\n",
    "        loss = recon_loss + kl_loss\n",
    "\n",
    "        # backward pass\n",
    "        loss.backward()\n",
    "        train_loss += loss.item()\n",
    "\n",
    "        # update the weights\n",
    "        optimizer.step()\n",
    "\n",
    "    return train_loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def test():\n",
    "    # set the evaluation mode\n",
    "    model.eval()\n",
    "\n",
    "    # test loss for the data\n",
    "    test_loss = 0\n",
    "\n",
    "    # we don't need to track the gradients, since we are not updating the parameters during evaluation / testing\n",
    "    with torch.no_grad():\n",
    "        for i, (x, _) in enumerate(test_iterator):\n",
    "            # reshape the data\n",
    "            x = x.view(-1, INPUT_DIM)\n",
    "            x = x.to(device)\n",
    "\n",
    "            # forward pass\n",
    "            x_sample, z_mu, z_var = model(x)\n",
    "\n",
    "            # reconstruction loss\n",
    "            recon_loss = F.binary_cross_entropy(x_sample, x, size_average=False)\n",
    "\n",
    "            # kl divergence loss\n",
    "            kl_loss = 0.5 * torch.sum(torch.exp(z_var) + z_mu**2 - 1.0 - z_var)\n",
    "\n",
    "            # total loss\n",
    "            loss = recon_loss + kl_loss\n",
    "            test_loss += loss.item()\n",
    "\n",
    "    return test_loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/envs/autolens/lib/python3.7/site-packages/torch/nn/_reduction.py:43: UserWarning: size_average and reduce args will be deprecated, please use reduction='sum' instead.\n",
      "  warnings.warn(warning.format(ret))\n"
     ]
    }
   ],
   "source": [
    "best_test_loss = float('inf')\n",
    "\n",
    "for e in range(N_EPOCHS):\n",
    "\n",
    "    train_loss = train()\n",
    "    test_loss = test()\n",
    "\n",
    "    train_loss /= len(train_dataset)\n",
    "    test_loss /= len(test_dataset)\n",
    "\n",
    "    print(f'Epoch {e}, Train Loss: {train_loss:.2f}, Test Loss: {test_loss:.2f}')\n",
    "\n",
    "    if best_test_loss > test_loss:\n",
    "        best_test_loss = test_loss\n",
    "        patience_counter = 1\n",
    "    else:\n",
    "        patience_counter += 1\n",
    "\n",
    "    if patience_counter > 3:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sample and generate a image\n",
    "z = torch.randn(1, LATENT_DIM).to(device)\n",
    "reconstructed_img = model.dec(z)\n",
    "img = reconstructed_img.view(150, 150).data\n",
    "\n",
    "print(z.shape)\n",
    "print(img.shape)\n",
    "plt.figure()\n",
    "plt.imshow(img, cmap='gray')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "\n",
    "# 3: CNN(incomplete)\n",
    "\n",
    "### Though it is already implemented in the paper(Deep Learning the Morphology of Dark Matter Substructure), this method is also useful when it comes to learning representation for anomaly detection."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TODO: make labels(or use directory names as labels) according to sub or no_sub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cnn_finetune import make_model\n",
    "\n",
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.datasets as datasets\n",
    "from torch.autograd import Variable\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "from PIL import Image\n",
    "import pandas as pd\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = 'pnasnet5large'\n",
    "CLASSES = ( 'sub', 'no_sub')\n",
    "BATCH_SIZE = 128\n",
    "MOMENTUM = 0.9 # for SDG\n",
    "N_EPOCHS = 10\n",
    "lr = 1e-3  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, dataloaders, criterion, optimizer, num_epochs=25, is_inception=False):\n",
    "    since = time.time()\n",
    "\n",
    "    val_acc_history = []\n",
    "\n",
    "    best_model_wts = copy.deepcopy(model.state_dict())\n",
    "    best_acc = 0.0\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        print('Epoch {}/{}'.format(epoch, num_epochs - 1))\n",
    "        print('-' * 10)\n",
    "\n",
    "        # Each epoch has a training and validation phase\n",
    "        for phase in ['train', 'val']:\n",
    "            if phase == 'train':\n",
    "                model.train()  # Set model to training mode\n",
    "            else:\n",
    "                model.eval()   # Set model to evaluate mode\n",
    "\n",
    "            running_loss = 0.0\n",
    "            running_corrects = 0\n",
    "\n",
    "            # Iterate over data.\n",
    "            for inputs, labels in dataloaders[phase]:\n",
    "                inputs = inputs.to(device)\n",
    "                labels = labels.to(device)\n",
    "\n",
    "                # zero the parameter gradients\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                # forward\n",
    "                # track history if only in train\n",
    "                with torch.set_grad_enabled(phase == 'train'):\n",
    "                    # Get model outputs and calculate loss\n",
    "                    # Special case for inception because in training it has an auxiliary output. In train\n",
    "                    #   mode we calculate the loss by summing the final output and the auxiliary output\n",
    "                    #   but in testing we only consider the final output.\n",
    "                    if is_inception and phase == 'train':\n",
    "                        # From https://discuss.pytorch.org/t/how-to-optimize-inception-model-with-auxiliary-classifiers/7958\n",
    "                        outputs, aux_outputs = model(inputs)\n",
    "                        loss1 = criterion(outputs, labels)\n",
    "                        loss2 = criterion(aux_outputs, labels)\n",
    "                        loss = loss1 + 0.4*loss2\n",
    "                    else:\n",
    "                        outputs = model(inputs)\n",
    "                        loss = criterion(outputs, labels)\n",
    "\n",
    "                    _, preds = torch.max(outputs, 1)\n",
    "\n",
    "                    # backward + optimize only if in training phase\n",
    "                    if phase == 'train':\n",
    "                        loss.backward()\n",
    "                        optimizer.step()\n",
    "\n",
    "                # statistics\n",
    "                running_loss += loss.item() * inputs.size(0)\n",
    "                running_corrects += torch.sum(preds == labels.data)\n",
    "\n",
    "            epoch_loss = running_loss / len(dataloaders[phase].dataset)\n",
    "            epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)\n",
    "\n",
    "            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))\n",
    "\n",
    "            # deep copy the model\n",
    "            if phase == 'val' and epoch_acc > best_acc:\n",
    "                best_acc = epoch_acc\n",
    "                best_model_wts = copy.deepcopy(model.state_dict())\n",
    "            if phase == 'val':\n",
    "                val_acc_history.append(epoch_acc)\n",
    "\n",
    "        print()\n",
    "\n",
    "    time_elapsed = time.time() - since\n",
    "    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))\n",
    "    print('Best val Acc: {:4f}'.format(best_acc))\n",
    "\n",
    "    # load best model weights\n",
    "    model.load_state_dict(best_model_wts)\n",
    "    return model, val_acc_history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = datasets.ImageFolder(root='./lenses/train', #3001 images each\n",
    "                                           transform=data_transform)\n",
    "test_dataset = datasets.ImageFolder(root='./lenses/test', #2000 images each\n",
    "                                           transform=data_transform)\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "        train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2\n",
    ")\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "        test_dataset, BATCH_SIZE, shuffle=False, num_workers=2\n",
    ")\n",
    "\n",
    "optimizer = optim.SGD(model.parameters(), lr=lr, momentum=MOMENTUM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model = make_model(\n",
    "#     MODEL_NAME, \n",
    "#     num_classes=len(CLASSES), \n",
    "#     pretrained=True, \n",
    "#     input_size=(150, 150)\n",
    "# ) \n",
    "# need to downgrade pytorch ver below 0.4 in order to use cnn_finetune\n",
    "\n",
    "model = torchvision.models.resnet18(pretrained=True)\n",
    "model.fc = nn.Linear(512, len(CLASSES))\n",
    "\n",
    "# model_ft = models.resnet18(pretrained=use_pretrained)\n",
    "#         set_parameter_requires_grad(model_ft, feature_extract)\n",
    "#         num_ftrs = model_ft.fc.in_features\n",
    "#         model_ft.fc = nn.Linear(num_ftrs, num_classes)\n",
    "#         input_size = 224\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create training and validation datasets\n",
    "image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ['train', 'test']}\n",
    "# Create training and validation dataloaders\n",
    "dataloaders_dict = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batch_size, shuffle=True, num_workers=4) for x in ['train', 'val']}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_ft = model.to(device)\n",
    "\n",
    "# Gather the parameters to be optimized/updated in this run. If we are\n",
    "#  finetuning we will be updating all parameters. However, if we are\n",
    "#  doing feature extract method, we will only update the parameters\n",
    "#  that we have just initialized, i.e. the parameters with requires_grad\n",
    "#  is True.\n",
    "params_to_update = model_ft.parameters()\n",
    "print(\"Params to learn:\")\n",
    "if feature_extract:\n",
    "    params_to_update = []\n",
    "    for name,param in model_ft.named_parameters():\n",
    "        if param.requires_grad == True:\n",
    "            params_to_update.append(param)\n",
    "            print(\"\\t\",name)\n",
    "else:\n",
    "    for name,param in model_ft.named_parameters():\n",
    "        if param.requires_grad == True:\n",
    "            print(\"\\t\",name)\n",
    "\n",
    "# Observe that all parameters are being optimized\n",
    "optimizer_ft = optim.SGD(params_to_update, lr=0.001, momentum=0.9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup the loss fxn\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "# Train and evaluate\n",
    "model_ft, hist = train_model(model_ft, dataloaders_dict, criterion, optimizer_ft, num_epochs=N_EPOCHS, is_inception=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for cnn_finetune\n",
    "# Use exponential decay for fine-tuning optimizer\n",
    "scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.975)\n",
    "\n",
    "# Train\n",
    "for epoch in range(1, N_EPOCHS + 1):\n",
    "    # Decay Learning Rate\n",
    "    scheduler.step(epoch)\n",
    "    train(model, epoch, optimizer, train_loader)\n",
    "    test(model, test_loader)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}