Noisy_Hook.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "6HiiAfjcTIpU"
      },
      "outputs": [],
      "source": [
        "%reset -f"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "8LJRXXN6gygC"
      },
      "outputs": [],
      "source": [
        "# !pip install brian2"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install torch"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wOq6G-QBejwM",
        "outputId": "17954e16-4bd4-4c54-c674-bd1064357d81"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (2.1.0+cu121)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch) (3.13.1)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from torch) (4.9.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch) (3.2.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch) (3.1.3)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch) (2023.6.0)\n",
            "Requirement already satisfied: triton==2.1.0 in /usr/local/lib/python3.10/dist-packages (from torch) (2.1.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch) (2.1.5)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch) (1.3.0)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install pytorch-lightning"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "czXuC1x8el5J",
        "outputId": "139cbe1b-b927-4ed2-9c7d-2117c6c3794f"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting pytorch-lightning\n",
            "  Downloading pytorch_lightning-2.1.4-py3-none-any.whl (778 kB)\n",
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/778.1 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[91m━━━━━━━━━━━━━\u001b[0m\u001b[90m╺\u001b[0m\u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m256.0/778.1 kB\u001b[0m \u001b[31m7.6 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m778.1/778.1 kB\u001b[0m \u001b[31m13.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: numpy>=1.17.2 in /usr/local/lib/python3.10/dist-packages (from pytorch-lightning) (1.23.5)\n",
            "Requirement already satisfied: torch>=1.12.0 in /usr/local/lib/python3.10/dist-packages (from pytorch-lightning) (2.1.0+cu121)\n",
            "Requirement already satisfied: tqdm>=4.57.0 in /usr/local/lib/python3.10/dist-packages (from pytorch-lightning) (4.66.1)\n",
            "Requirement already satisfied: PyYAML>=5.4 in /usr/local/lib/python3.10/dist-packages (from pytorch-lightning) (6.0.1)\n",
            "Requirement already satisfied: fsspec[http]>=2022.5.0 in /usr/local/lib/python3.10/dist-packages (from pytorch-lightning) (2023.6.0)\n",
            "Collecting torchmetrics>=0.7.0 (from pytorch-lightning)\n",
            "  Downloading torchmetrics-1.3.0.post0-py3-none-any.whl (840 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m840.2/840.2 kB\u001b[0m \u001b[31m27.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from pytorch-lightning) (23.2)\n",
            "Requirement already satisfied: typing-extensions>=4.0.0 in /usr/local/lib/python3.10/dist-packages (from pytorch-lightning) (4.9.0)\n",
            "Collecting lightning-utilities>=0.8.0 (from pytorch-lightning)\n",
            "  Downloading lightning_utilities-0.10.1-py3-none-any.whl (24 kB)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from fsspec[http]>=2022.5.0->pytorch-lightning) (2.31.0)\n",
            "Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1 in /usr/local/lib/python3.10/dist-packages (from fsspec[http]>=2022.5.0->pytorch-lightning) (3.9.3)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from lightning-utilities>=0.8.0->pytorch-lightning) (67.7.2)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch>=1.12.0->pytorch-lightning) (3.13.1)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=1.12.0->pytorch-lightning) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.12.0->pytorch-lightning) (3.2.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.12.0->pytorch-lightning) (3.1.3)\n",
            "Requirement already satisfied: triton==2.1.0 in /usr/local/lib/python3.10/dist-packages (from torch>=1.12.0->pytorch-lightning) (2.1.0)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2022.5.0->pytorch-lightning) (1.3.1)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2022.5.0->pytorch-lightning) (23.2.0)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2022.5.0->pytorch-lightning) (1.4.1)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2022.5.0->pytorch-lightning) (6.0.5)\n",
            "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2022.5.0->pytorch-lightning) (1.9.4)\n",
            "Requirement already satisfied: async-timeout<5.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2022.5.0->pytorch-lightning) (4.0.3)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.12.0->pytorch-lightning) (2.1.5)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->fsspec[http]>=2022.5.0->pytorch-lightning) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->fsspec[http]>=2022.5.0->pytorch-lightning) (3.6)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->fsspec[http]>=2022.5.0->pytorch-lightning) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->fsspec[http]>=2022.5.0->pytorch-lightning) (2024.2.2)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=1.12.0->pytorch-lightning) (1.3.0)\n",
            "Installing collected packages: lightning-utilities, torchmetrics, pytorch-lightning\n",
            "Successfully installed lightning-utilities-0.10.1 pytorch-lightning-2.1.4 torchmetrics-1.3.0.post0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MTEDV8wqLZXr",
        "outputId": "116776a8-c81c-469d-c5bb-5cd5bdd73a21"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting snntorch\n",
            "  Downloading snntorch-0.7.0-py2.py3-none-any.whl (108 kB)\n",
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/109.0 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[91m╸\u001b[0m\u001b[90m━━\u001b[0m \u001b[32m102.4/109.0 kB\u001b[0m \u001b[31m3.2 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m109.0/109.0 kB\u001b[0m \u001b[31m2.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: torch>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from snntorch) (2.1.0+cu121)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.10/dist-packages (from snntorch) (1.5.3)\n",
            "Requirement already satisfied: matplotlib in /usr/local/lib/python3.10/dist-packages (from snntorch) (3.7.1)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from snntorch) (1.23.5)\n",
            "Collecting nir (from snntorch)\n",
            "  Downloading nir-1.0.1-py3-none-any.whl (76 kB)\n",
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/76.2 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m76.2/76.2 kB\u001b[0m \u001b[31m6.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting nirtorch (from snntorch)\n",
            "  Downloading nirtorch-1.0-py3-none-any.whl (13 kB)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch>=1.1.0->snntorch) (3.13.1)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from torch>=1.1.0->snntorch) (4.9.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=1.1.0->snntorch) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.1.0->snntorch) (3.2.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.1.0->snntorch) (3.1.3)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch>=1.1.0->snntorch) (2023.6.0)\n",
            "Requirement already satisfied: triton==2.1.0 in /usr/local/lib/python3.10/dist-packages (from torch>=1.1.0->snntorch) (2.1.0)\n",
            "Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->snntorch) (1.2.0)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.10/dist-packages (from matplotlib->snntorch) (0.12.1)\n",
            "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib->snntorch) (4.47.2)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->snntorch) (1.4.5)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib->snntorch) (23.2)\n",
            "Requirement already satisfied: pillow>=6.2.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib->snntorch) (9.4.0)\n",
            "Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->snntorch) (3.1.1)\n",
            "Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.10/dist-packages (from matplotlib->snntorch) (2.8.2)\n",
            "Requirement already satisfied: h5py in /usr/local/lib/python3.10/dist-packages (from nir->snntorch) (3.9.0)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas->snntorch) (2023.4)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.7->matplotlib->snntorch) (1.16.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.1.0->snntorch) (2.1.5)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=1.1.0->snntorch) (1.3.0)\n",
            "Installing collected packages: nir, nirtorch, snntorch\n",
            "Successfully installed nir-1.0.1 nirtorch-1.0 snntorch-0.7.0\n"
          ]
        }
      ],
      "source": [
        "!pip install snntorch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NGPDVXvjdGsl",
        "outputId": "aed1cd9f-b197-4e63-c8de-413d56d5b752"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: torchsummary in /usr/local/lib/python3.10/dist-packages (1.5.1)\n"
          ]
        }
      ],
      "source": [
        "!pip install torchsummary"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "V7xJifRLMc2Q"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KehrXuVuPiAt"
      },
      "source": [
        "## set seeds for PyTorch and Numpy to ensure reproducibility:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "NbSadcLWPedr"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import random\n",
        "import numpy as np\n",
        "\n",
        "# Set seeds for Python, Numpy, and Torch for reproducibility\n",
        "seed = 42\n",
        "random.seed(seed)\n",
        "np.random.seed(seed)\n",
        "torch.manual_seed(seed)\n",
        "torch.cuda.manual_seed_all(seed)\n",
        "\n",
        "# Additional steps if you're using GPU\n",
        "torch.backends.cudnn.deterministic = True\n",
        "torch.backends.cudnn.benchmark = False\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "UyXSF4IqLFND"
      },
      "outputs": [],
      "source": [
        "# imports\n",
        "import pickle\n",
        "import matplotlib.animation as animation\n",
        "from scipy.integrate import simps\n",
        "import torch\n",
        "torch.cuda.empty_cache()\n",
        "import os, sys, time, datetime, json, random\n",
        "import snntorch as snn\n",
        "from snntorch import spikeplot as splt\n",
        "from snntorch import spikegen\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import matplotlib.pyplot as plt\n",
        "import os\n",
        "import torch.nn.functional as F\n",
        "from torchvision import datasets, transforms\n",
        "from torch.utils.data import DataLoader\n",
        "from torchvision import utils as utls\n",
        "from snntorch import utils\n",
        "from snntorch import surrogate\n",
        "import numpy as np\n",
        "import math\n",
        "from sklearn.metrics import auc\n",
        "from torchsummary import summary\n",
        "from sklearn.manifold import TSNE\n",
        "from sklearn.decomposition import PCA\n",
        "import torch.nn as nn\n",
        "from scipy.stats import entropy\n",
        "from scipy.special import kl_div\n",
        "from torch.autograd import Variable\n",
        "# import spikeflow as snn"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# torch.autograd.set_detect_anomaly(True)"
      ],
      "metadata": {
        "id": "ZDOWUXysMxFD"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "nOyTY0dE50vG",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0aa2c326-cb79-4dcd-b42a-04dc90fd2d17"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\n",
            "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to dataset/MNIST/raw/train-images-idx3-ubyte.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 9912422/9912422 [00:00<00:00, 111665707.58it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting dataset/MNIST/raw/train-images-idx3-ubyte.gz to dataset/MNIST/raw\n",
            "\n",
            "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz\n",
            "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to dataset/MNIST/raw/train-labels-idx1-ubyte.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28881/28881 [00:00<00:00, 41756530.10it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting dataset/MNIST/raw/train-labels-idx1-ubyte.gz to dataset/MNIST/raw\n",
            "\n",
            "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz\n",
            "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to dataset/MNIST/raw/t10k-images-idx3-ubyte.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 1648877/1648877 [00:00<00:00, 31887660.15it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting dataset/MNIST/raw/t10k-images-idx3-ubyte.gz to dataset/MNIST/raw\n",
            "\n",
            "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\n",
            "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to dataset/MNIST/raw/t10k-labels-idx1-ubyte.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 4542/4542 [00:00<00:00, 17688513.25it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting dataset/MNIST/raw/t10k-labels-idx1-ubyte.gz to dataset/MNIST/raw\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# /////////////////////# Building the Autoencoder\n",
        "#-------------------DataLoaders.  using the MNIST dataset\n",
        "\n",
        "# dataloader arguments\n",
        "batch_size = 250\n",
        "data_path='/data/mnist'\n",
        "\n",
        "dtype = torch.float\n",
        "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
        "\n",
        "\n",
        "# /////////////////////////////////# Define a transform\n",
        "input_size = 32 # resizing the original MNIST from 28 to 32\n",
        "\n",
        "transform = transforms.Compose([\n",
        "            transforms.Resize((input_size, input_size)),\n",
        "            transforms.Grayscale(),\n",
        "            transforms.ToTensor(),\n",
        "            transforms.Normalize((0,), (1,))])\n",
        "\n",
        "#------------------------------------------- Load MNIST\n",
        "# Training data\n",
        "train_dataset = datasets.MNIST(root='dataset/', train=True, transform=transform, download=True)\n",
        "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
        "\n",
        "# Testing data\n",
        "test_dataset = datasets.MNIST(root='dataset/', train=False, transform=transform, download=True)\n",
        "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "Ha5yn2VjO39Q"
      },
      "outputs": [],
      "source": [
        "# creating directories where we can save the original and reconstructed images for training and testing:\n",
        "# create training/ and testing/ folders in the chosen path\n",
        "if not os.path.isdir('figures/training'):\n",
        "    os.makedirs('figures/training')\n",
        "if not os.path.isdir('figures/binarytraining'):\n",
        "    os.makedirs('figures/binarytraining')\n",
        "\n",
        "if not os.path.isdir('figures/testing'):\n",
        "    os.makedirs('figures/testing')\n",
        "if not os.path.isdir('figures/binarytesting'):\n",
        "    os.makedirs('figures/binarytesting')\n",
        "\n",
        "\n",
        "if not os.path.isdir('Saved_Trained_Checkpoints/'):\n",
        "    os.makedirs('Saved_Trained_Checkpoints/')\n",
        "\n",
        "if not os.path.isdir('Output_Spikes/'):\n",
        "    os.makedirs('Output_Spikes/')\n",
        "\n",
        "if not os.path.isdir('Enc_syn_Spikes/'):\n",
        "    os.makedirs('Enc_syn_Spikes/')\n",
        "\n",
        "\n",
        "if not os.path.isdir('Intermediate_Lyrs/'):\n",
        "    os.makedirs('Intermediate_Lyrs/')\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CcZMZrCoKFaq",
        "outputId": "bb91669c-b29e-4c56-8cc3-b3ef42566b43"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Using CPU\n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "\n",
        "if torch.cuda.is_available():\n",
        "    device = torch.device(\"cuda\")\n",
        "    print(f\"Using {torch.cuda.get_device_name()} ({device})\")\n",
        "else:\n",
        "    device = torch.device(\"cpu\")\n",
        "    print(\"Using CPU\")\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "_IMynIKMTLJV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kyQvxnL5umDu"
      },
      "source": [
        "### To manipulate and test (Working for Encoder and Decoder outputs):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2VQHpS2xufu0"
      },
      "outputs": [],
      "source": [
        "# class SAE(nn.Module):\n",
        "#     def __init__(self):\n",
        "#         super().__init__()\n",
        "\n",
        "\n",
        "#         # Encoder\n",
        "#         self.encoder = nn.Sequential(\n",
        "#                             nn.Conv2d(1, 32, 3,padding = 1,stride=2), # Conv Layer 1\n",
        "#                             # torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1,\n",
        "#                                 # bias=True, padding_mode='zeros',  device=None, dtype=None)\n",
        "#                             nn.BatchNorm2d(32),\n",
        "#                             snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "#                             # snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True,threshold=thresh), #SNN TORCH LIF NEURON\n",
        "#                             nn.Conv2d(32, 64, 3,padding = 1,stride=2), # Conv Layer 2\n",
        "#                             nn.BatchNorm2d(64),\n",
        "#                             # snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "#                             snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True,threshold=thresh), #SNN TORCH LIF NEURON\n",
        "#                             # snn.Alpha(alpha=alpha1, beta=beta1, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "#                             nn.Conv2d(64, 128, 3,padding = 1,stride=2), # Conv Layer 3\n",
        "#                             nn.BatchNorm2d(128),\n",
        "#                             # snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "#                             snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True,threshold=thresh), #SNN TORCH LIF NEURON\n",
        "#                             # snn.Alpha(alpha=alpha11, beta=beta11, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "#                             nn.Flatten(start_dim = 1, end_dim = 3), #Flatten convolutional output\n",
        "#                             nn.Linear(128*4*4, latent_dim), # Fully connected linear layer\n",
        "#                             snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True,threshold=thresh)\n",
        "#                             # snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True, output=True,threshold=thresh)\n",
        "#                             )\n",
        "\n",
        "\n",
        "#         self.latent_dim = latent_dim #dimensions of the encoded z-space data\n",
        "\n",
        "#         #ve from the flattened encoded representation (latent_dim) back to a tensor representation to\n",
        "#             # use in transposed convolution.\n",
        "#           # To do so, we need to run an additional fully-connected linear layer transforming the data back into a tensor of 128 x 4 x 4:\n",
        "\n",
        "#         self.linearNet= nn.Sequential(\n",
        "#                                       nn.Linear(latent_dim,128*4*4),\n",
        "#                                       snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True,threshold=thresh))\n",
        "#                                       #snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True, output=True,threshold=thresh))\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "#       #  The decoder, with three transposed convolutional (nn.ConvTranspose2d) layers and one linear output layer.\n",
        "#       # Although we converted the latent data back into tensor form for convolution, we still need to Unflatten it to a tensor of 128 x 4 x 4,\n",
        "#         # as the input to the network is 1 dimensional.  This is done using nn.Unflatten in the first line of the Decoder:\n",
        "#         # Decoder:\n",
        "#         self.decoder = nn.Sequential(\n",
        "#                             nn.Unflatten(1,(128,4,4)), #Unflatten data from 1 dim to tensor of 128 x 4 x 4\n",
        "#                             snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "#                             # snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "#                             nn.ConvTranspose2d(128, 64, 3,padding = 1,stride=(2,2),output_padding=1),\n",
        "#                             nn.BatchNorm2d(64),\n",
        "#                             # snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "#                             snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "#                             # snn.Alpha(alpha=alpha2, beta=beta2, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "#                             nn.ConvTranspose2d(64, 32, 3,padding = 1,stride=(2,2),output_padding=1),\n",
        "#                             nn.BatchNorm2d(32),\n",
        "#                             # snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "#                             snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "#                             # snn.Alpha(alpha=alpha22, beta=beta22, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "#                             nn.ConvTranspose2d(32, 1, 3,padding = 1,stride=(2,2),output_padding=1),\n",
        "#                             snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,output=True,threshold=20000) #so membrane can be trained\n",
        "#                             # snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,output=True,threshold=thresh) #---------------------------------------------- ADDED\n",
        "#                             # snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True,output=True,threshold=20000) #so membrane can be trained\n",
        "#                             )\n",
        "#         # final Leaky layer, our spiking threshold (thresh) is set extremely high. This is a neat trick in snnTorch, which allows the neuron\n",
        "#         # membrane in the final layer to continuously be updated, without ever reaching a spiking threshold.\n",
        "\n",
        "#         # using the membrane potential output from the final layer for the image reconstruction.\n",
        "#             # snnTorch allows us to use either the spikes or membrane potential of each neuron in training.\n",
        "\n",
        "\n",
        "#     def forward(self, x):\n",
        "#         utils.reset(self.encoder) #need to reset the hidden states of LIF\n",
        "#         utils.reset(self.decoder)\n",
        "#         utils.reset(self.linearNet)\n",
        "\n",
        "#     #-----------------------------encode\n",
        "#         spk_mem=[];\n",
        "#         spk_rec=[];\n",
        "#         spk_rec_syn=[];\n",
        "#         encoder_mem=[];\n",
        "#         spk_rec_dec=[];\n",
        "#         spk_mem_dec=[];\n",
        "#         for step in range(num_steps): #for t in time\n",
        "#             spk_x, mem_x = self.encode(x) #Output spike trains and neuron membrane states\n",
        "#             # print(\"Size of = self.encode(x):\", self.encode(x).size())\n",
        "#             # print(\"len of = self.encode(x):\", len(self.encode(x)))\n",
        "#             spk_rec.append(spk_x)\n",
        "#             spk_mem.append(mem_x)\n",
        "\n",
        "#         spk_rec=torch.stack(spk_rec,dim=2) # stack spikes in second tensor dimension\n",
        "#         spk_mem=torch.stack(spk_mem,dim=2) # stack membranes in second tensor dimension\n",
        "# # torch.stack joins (concatenates) a sequence of tensors (two or more tensors) along a new dimension.\n",
        "#         # print(\"Size of spk_rec:\", spk_rec.size())   # Size of spk_rec: torch.Size([250, 32, 5])\n",
        "#         # print(\"Size of spk_mem:\", spk_mem.size()) # Size of spk_mem: torch.Size([250, 32, 5])\n",
        "#         # out_en = spk_rec[:,:,-1]     #//////////////////////////////////////ADDED---------------- shows  batch size (250) different examples on (32) channel number\n",
        "#         out_en = spk_rec[0, :,:] #//////////////////////////////////////ADDED------------ latent dim (32) and time (num_steps)(5)\n",
        "#         # print(\"Size of out_en:\", out_en.size()) #Size of out_en: torch.Size([250, 32])\n",
        "\n",
        "#     #------------------------------decode\n",
        "#         spk_mem2=[];\n",
        "#         spk_rec2=[];\n",
        "#         decoded_x=[];\n",
        "#         spk_x_dec=[];\n",
        "#         mem_x_dec=[];\n",
        "#         for step in range(num_steps): #for t in time\n",
        "#             # spk_x_dec, mem_x_dec = self.decode(spk_rec[...,step]) #//////////////////////////ADDED\n",
        "#             x_recon, x_mem_recon = self.decode(spk_rec[...,step])\n",
        "\n",
        "#             spk_rec2.append(x_recon)\n",
        "#             spk_mem2.append(x_mem_recon)\n",
        "\n",
        "#             # spk_rec_dec.append(spk_x_dec)   #//////////////////////////ADDED Size of spk_rec_dec: torch.Size([250, 1, 5, 32, 32])\n",
        "#             # spk_mem_dec.append(mem_x_dec)   #//////////////////////////ADDED\n",
        "\n",
        "#         # spk_rec_dec=torch.stack(spk_rec_dec,dim=2) #//////////////////////////ADDED\n",
        "#         # spk_mem_dec=torch.stack(spk_mem_dec,dim=2) #//////////////////////////ADDED\n",
        "#         spk_rec2=torch.stack(spk_rec2,dim=4)\n",
        "#         spk_mem2=torch.stack(spk_mem2,dim=4)\n",
        "\n",
        "#         # out = spk_mem2[:,:,:,:,-1] #return the membrane potential of the output neuron at t = -1 (last t)\n",
        "#         # out = spk_rec2[:,:,:,:,-1]  #//////////////////////////ADDED #  Size of out: torch.Size([250, 1, 32, 32])\n",
        "#         # out = spk_rec_dec[:,:,-1] #//////////////////////////ADDED  # Size of out: torch.Size([250, 1, 32, 32])\n",
        "#         # out = spk_rec_dec[:,0,:,-1] #//////////////////////////ADDED  # Size of out: torch.Size([250, 5, 32])\n",
        "#         # out = spk_rec_dec[:,0,0,:,-1] #//////////////////////////ADDED  # Size of out: torch.Size([250, 32])\n",
        "#         # out = spk_rec_dec[:,-1] #//////////////////////////ADDED  Size of out: torch.Size([250, 5, 32, 32])\n",
        "#         # out = spk_rec2[0,0,:,:] #//////////////////////////ADDED Now working with [250, 1, 32, 32]) because it is (torch.Size([32, 32, 5]))\n",
        "#         out = spk_mem2[:,:,:,:,-1]  #//////////////////////////ADDED #  Size of out: torch.Size([250, 1, 32, 32])--------------- for one digit at the last time t\n",
        "\n",
        "#         # print(\"Size of out:\", out.size())\n",
        "\n",
        "#         # Save the out_en tensor\n",
        "#         self.out_en = out_en\n",
        "#         self.out = out\n",
        "\n",
        "#         return out, out_en\n",
        "\n",
        "\n",
        "#     def encode(self,x):\n",
        "#       spk_latent_x, mem_latent_x = self.encoder(x)\n",
        "#       return spk_latent_x, mem_latent_x, syn1, memsyn1\n",
        "\n",
        "#     def decode(self,x):\n",
        "#         spk_x, mem_x = self.linearNet(x) #convert latent dimension back to total size of features in encoder final layer\n",
        "#         spk_x2, mem_x2 = self.decoder(spk_x)\n",
        "#         return spk_x2, mem_x2"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Working with \"out\" and \"out_en\":\n"
      ],
      "metadata": {
        "id": "vIXf5IFdcTt9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# class SAE(nn.Module):\n",
        "#     def __init__(self):\n",
        "#         super().__init__()\n",
        "#         # Encoder\n",
        "#         self.encoder = nn.Sequential(\n",
        "#                             nn.Conv2d(1, 32, 3,padding = 1,stride=2), # Conv Layer 1\n",
        "#                             nn.BatchNorm2d(32),\n",
        "#                             snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "#                             nn.Conv2d(32, 64, 3,padding = 1,stride=2), # Conv Layer 2\n",
        "#                             nn.BatchNorm2d(64),\n",
        "#                             snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True,threshold=thresh), #SNN TORCH LIF NEURON\n",
        "#                             nn.Conv2d(64, 128, 3,padding = 1,stride=2), # Conv Layer 3\n",
        "#                             nn.BatchNorm2d(128),\n",
        "#                             snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True,threshold=thresh), #SNN TORCH LIF NEURON\n",
        "#                             nn.Flatten(start_dim = 1, end_dim = 3), #Flatten convolutional output\n",
        "#                             nn.Linear(128*4*4, latent_dim), # Fully connected linear layer\n",
        "#                             snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True,threshold=thresh)\n",
        "#                             )\n",
        "#         self.latent_dim = latent_dim #dimensions of the encoded z-space data\n",
        "#         self.linearNet= nn.Sequential(\n",
        "#                                       nn.Linear(latent_dim,128*4*4),\n",
        "#                                       snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True,threshold=thresh))\n",
        "\n",
        "#         # Decoder:\n",
        "#         self.decoder = nn.Sequential(\n",
        "#                             nn.Unflatten(1,(128,4,4)), #Unflatten data from 1 dim to tensor of 128 x 4 x 4\n",
        "#                             snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "#                             nn.ConvTranspose2d(128, 64, 3,padding = 1,stride=(2,2),output_padding=1),\n",
        "#                             nn.BatchNorm2d(64),\n",
        "#                             snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "#                             nn.ConvTranspose2d(64, 32, 3,padding = 1,stride=(2,2),output_padding=1),\n",
        "#                             nn.BatchNorm2d(32),\n",
        "#                             snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "#                             nn.ConvTranspose2d(32, 1, 3,padding = 1,stride=(2,2),output_padding=1),\n",
        "#                             snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,output=True,threshold=20000) #so membrane can be trained\n",
        "#                             )\n",
        "#     def forward(self, x):\n",
        "#         utils.reset(self.encoder) #need to reset the hidden states of LIF\n",
        "#         utils.reset(self.decoder)\n",
        "#         utils.reset(self.linearNet)\n",
        "\n",
        "#     #-----------------------------encode\n",
        "#         spk_mem=[];\n",
        "#         spk_rec=[];\n",
        "#         spk_rec_syn=[];\n",
        "#         encoder_mem=[];\n",
        "#         spk_rec_dec=[];\n",
        "#         spk_mem_dec=[];\n",
        "#         for step in range(num_steps): #for t in time\n",
        "#             spk_x, mem_x = self.encode(x) #Output spike trains and neuron membrane states\n",
        "#             # print(\"spk_x in spk_x, mem_x = self.encode(x): \",spk_x.size())    #                       torch.Size([250, 32])\n",
        "#             # print(\"mem_x in spk_x, mem_x = self.encode(x): \",mem_x.size())    #                       torch.Size([250, 32])\n",
        "#             spk_rec.append(spk_x)\n",
        "#             spk_mem.append(mem_x)\n",
        "#             # print(\"len of spk_rec in spk_rec.append(spk_x): \",len(spk_rec))\n",
        "#             # print(\"len of spk_mem in spk_rec.append(mem_x): \",len(spk_mem))\n",
        "\n",
        "#         spk_rec=torch.stack(spk_rec,dim=2) # stack spikes in second tensor dimension # ----------------spk_rec in torch.stack(spk_rec,dim=2):  torch.Size([250, 32, 5])\n",
        "#         # print(\"----------------spk_rec in torch.stack(spk_rec,dim=2): \",spk_rec.size())\n",
        "#         spk_mem=torch.stack(spk_mem,dim=2) # stack membranes in second tensor dimension # ----------------spk_mem in torch.stack(spk_mem,dim=2):  torch.Size([250, 32, 5])\n",
        "#         # print(\"----------------spk_mem in torch.stack(spk_mem,dim=2): \",spk_mem.size())\n",
        "#         # out_en = spk_rec[0, :,:] #//////////////////////////////////////ADDED------------ latent dim (32) and time (num_steps)(5) # out_en in spk_rec[0, :,:]:----------- torch.Size([32, 5])\n",
        "#         # print(\"out_en in spk_rec[0, :,:]:-----------\" , out_en.size())\n",
        "#         out_en = spk_rec[...,step]\n",
        "#         # print(\"spk_rec[...,step]:-----------\" , spk_rec[...,step].size()) # spk_rec[...,step]:----------- torch.Size([250, 32])       input of the latent and then decoder\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "#     #------------------------------decode\n",
        "#         spk_mem2=[];\n",
        "#         spk_rec2=[];\n",
        "#         decoded_x=[];\n",
        "#         spk_x_dec=[];\n",
        "#         mem_x_dec=[];\n",
        "#         for step in range(num_steps): #for t in time                                   from decoder: ([250, 1, 32, 32])\n",
        "#             x_recon, x_mem_recon = self.decode(spk_rec[...,step])                             # x_recon in self.decode(spk_rec[...,step]): torch.Size([250, 1, 32, 32])\n",
        "#             # print(\"x_recon in self.decode(spk_rec[...,step]):\" , x_recon.size())            # x_mem_recon in self.decode(spk_rec[...,step]): torch.Size([250, 1, 32, 32])\n",
        "#             # print(\"x_mem_recon in self.decode(spk_rec[...,step]):\" , x_mem_recon.size())\n",
        "\n",
        "#             spk_rec2.append(x_recon)\n",
        "#             spk_mem2.append(x_mem_recon)\n",
        "#             # print(\"len of spk_rec2.append(x_recon)\",len(spk_rec2))\n",
        "#             # print(\"len of spk_mem2.append(x_recon)\",len(spk_mem2))\n",
        "\n",
        "#         spk_rec2=torch.stack(spk_rec2,dim=4)\n",
        "#         spk_mem2=torch.stack(spk_mem2,dim=4)\n",
        "#         # print(\"spk_rec2 in torch.stack(spk_rec2,dim=4):\", spk_rec2.size())          # spk_rec2 in torch.stack(spk_rec2,dim=4): torch.Size([250, 1, 32, 32, 5])\n",
        "#         # print(\"spk_mem2 in torch.stack(spk_rec2,dim=4):\", spk_mem2.size())          # spk_mem2 in torch.stack(spk_rec2,dim=4): torch.Size([250, 1, 32, 32, 5])\n",
        "\n",
        "#         out = spk_mem2[:,:,:,:,-1]  #//////////////////////////ADDED #  Size of out: torch.Size([250, 1, 32, 32])--------------- for one digit at the last time t\n",
        "\n",
        "#         # self.out_en = out_en\n",
        "#         # self.out = out\n",
        "\n",
        "#         return out, out_en\n",
        "\n",
        "\n",
        "#     def encode(self,x):\n",
        "#       spk_latent_x, mem_latent_x = self.encoder(x)\n",
        "#       return spk_latent_x, mem_latent_x\n",
        "\n",
        "#     def decode(self,x):\n",
        "#         spk_x, mem_x = self.linearNet(x) #convert latent dimension back to total size of features in encoder final layer\n",
        "#         spk_x2, mem_x2 = self.decoder(spk_x)\n",
        "#         return spk_x2, mem_x2"
      ],
      "metadata": {
        "id": "MuOth1ZabvgL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class SAE(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "\n",
        "\n",
        "        # encoder_layers = [\n",
        "        #     ('conv1', nn.Conv2d(1, 32, 3, padding=1, stride=2)),  # Conv Layer 1\n",
        "        #     ('batchnorm1', nn.BatchNorm2d(32)),\n",
        "        #     ('leaky1', snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True, threshold=thresh)),\n",
        "        #     ('conv2', nn.Conv2d(32, 64, 3, padding=1, stride=2)),  # Conv Layer 2\n",
        "        #     ('batchnorm2', nn.BatchNorm2d(64)),\n",
        "        #     ('synaptic1', snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True, output=True, threshold=thresh)),  # SNN TORCH LIF NEURON\n",
        "        #     ('conv3', nn.Conv2d(64, 128, 3, padding=1, stride=2)),  # Conv Layer 3\n",
        "        #     ('batchnorm3', nn.BatchNorm2d(128)),\n",
        "        #     ('synaptic2', snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True, output=True, threshold=thresh)),  # SNN TORCH LIF NEURON\n",
        "        #     ('flatten', nn.Flatten(start_dim=1, end_dim=3)),  # Flatten convolutional output\n",
        "        #     ('linear', nn.Linear(128 * 4 * 4, latent_dim)),  # Fully connected linear layer\n",
        "        #     ('leaky2', snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True, threshold=thresh))\n",
        "        # ]\n",
        "\n",
        "\n",
        "        # Encoder\n",
        "        self.encoder = nn.Sequential(\n",
        "                            nn.Conv2d(1, 32, 3,padding = 1,stride=2), # Conv Layer 1\n",
        "                            nn.BatchNorm2d(32),\n",
        "                            snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "                            nn.Conv2d(32, 64, 3,padding = 1,stride=2), # Conv Layer 2\n",
        "                            nn.BatchNorm2d(64),\n",
        "                            snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True, threshold=thresh), #SNN TORCH LIF NEURON\n",
        "                            nn.Conv2d(64, 128, 3,padding = 1,stride=2), # Conv Layer 3\n",
        "                            nn.BatchNorm2d(128),\n",
        "                            snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True,threshold=thresh), #SNN TORCH LIF NEURON\n",
        "                            nn.Flatten(start_dim = 1, end_dim = 3), #Flatten convolutional output\n",
        "                            nn.Linear(128*4*4, latent_dim), # Fully connected linear layer\n",
        "                            snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True,threshold=thresh)\n",
        "                            )\n",
        "\n",
        "\n",
        "        self.latent_dim = latent_dim #dimensions of the encoded z-space data\n",
        "        self.linearNet= nn.Sequential(\n",
        "                                      nn.Linear(latent_dim,128*4*4),\n",
        "                                      snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True,threshold=thresh))\n",
        "\n",
        "        # Decoder:\n",
        "        self.decoder = nn.Sequential(\n",
        "                            nn.Unflatten(1,(128,4,4)), #Unflatten data from 1 dim to tensor of 128 x 4 x 4\n",
        "                            snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "                            nn.ConvTranspose2d(128, 64, 3,padding = 1,stride=(2,2),output_padding=1),\n",
        "                            nn.BatchNorm2d(64),\n",
        "                            snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "                            nn.ConvTranspose2d(64, 32, 3,padding = 1,stride=(2,2),output_padding=1),\n",
        "                            nn.BatchNorm2d(32),\n",
        "                            snn.Synaptic(alpha=alpha, beta=beta_syn, spike_grad=spike_grad, init_hidden=True,threshold=thresh),\n",
        "                            nn.ConvTranspose2d(32, 1, 3,padding = 1,stride=(2,2),output_padding=1),\n",
        "                            snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,output=True,threshold=20000) #so membrane can be trained\n",
        "                            )\n",
        "    def forward(self, x):\n",
        "        utils.reset(self.encoder) #need to reset the hidden states of LIF\n",
        "        utils.reset(self.decoder)\n",
        "        utils.reset(self.linearNet)\n",
        "\n",
        "    #-----------------------------encode\n",
        "        spk_mem=[];\n",
        "        spk_rec=[];\n",
        "        spk_rec_syn=[];\n",
        "        encoder_mem=[];\n",
        "        spk_rec_dec=[];\n",
        "        spk_mem_dec=[];\n",
        "        enc5_rec = [];\n",
        "\n",
        "\n",
        "     #------------------------------ intermediate layers\n",
        "\n",
        "        # for step in range(num_steps):\n",
        "        #     enc5 = self.encoder[5](x)             #  enc5 shape: torch.Size([250, 1, 32, 32])\n",
        "        #     enc5_rec.append(enc5)\n",
        "        # enc5_rec = torch.stack(enc5_rec, dim=2)            #   enc5_rec size: torch.Size([250, 1, 5, 32, 32])\n",
        "        # Enc_syn_1 = enc5_rec[:, :, -1]                      # #   torch.Size([250, 1, 32, 32])\n",
        "\n",
        "     #------------------------------ encode\n",
        "        for step in range(num_steps):\n",
        "            spk_x, mem_x = self.encoder(x)              # spk_x size: ([250, 32])  ,   mem_x size: ([250, 32])  , x.shape : torch.Size([250, 1, 32, 32])\n",
        "            spk_rec.append(spk_x)\n",
        "            spk_mem.append(mem_x)\n",
        "\n",
        "        spk_rec=torch.stack(spk_rec,dim=2) # stack spikes in second tensor dimension # ----------------spk_rec in torch.stack(spk_rec,dim=2):  torch.Size([250, 32, 5])\n",
        "        spk_mem=torch.stack(spk_mem,dim=2) # stack membranes in second tensor dimension # ----------------spk_mem in torch.stack(spk_mem,dim=2):  torch.Size([250, 32, 5])\n",
        "        out_en = spk_rec[...,step]\n",
        "\n",
        "#ADDED---------------------------------------------\n",
        "        # Apply Poisson noise to spike trains\n",
        "        noise_level = 0.1  # Adjust the noise level\n",
        "        noise = torch.poisson(out_en * noise_level)\n",
        "        # Add noise to the signal\n",
        "        noisy_spk_rec = out_en + noise\n",
        "\n",
        "# -------------------------------------------------\n",
        "        # print(\"out_en= spk_rec[...,step]:-----------\" , spk_rec[...,step].size()) # spk_rec[...,step]:----------- torch.Size([250, 32])       input of the latent and then decoder\n",
        "\n",
        "     #------------------------------decode\n",
        "        spk_mem2=[];\n",
        "        spk_rec2=[];\n",
        "        decoded_x=[];\n",
        "        spk_x_dec=[];\n",
        "        mem_x_dec=[];\n",
        "        for step in range(num_steps): #for t in time                           #        from decoder: ([250, 1, 32, 32])\n",
        "            x_recon, x_mem_recon = self.decode(spk_rec[...,step])\n",
        "            spk_rec2.append(x_recon)\n",
        "            spk_mem2.append(x_mem_recon)\n",
        "\n",
        "        spk_rec2=torch.stack(spk_rec2,dim=4)\n",
        "        spk_mem2=torch.stack(spk_mem2,dim=4)\n",
        "\n",
        "        out = spk_mem2[:,:,:,:,-1]\n",
        "\n",
        "        self.out_en = out_en\n",
        "        self.out = out\n",
        "\n",
        "        return out, out_en\n",
        "        # return out, Enc_syn_1\n",
        "\n",
        "    def encode(self,x):\n",
        "      spk_latent_x, mem_latent_x = self.encoder(x)\n",
        "      return spk_latent_x, mem_latent_x\n",
        "\n",
        "\n",
        "    def decode(self,x):\n",
        "        spk_x, mem_x = self.linearNet(x) #convert latent dimension back to total size of features in encoder final layer\n",
        "        spk_x2, mem_x2 = self.decoder(spk_x)\n",
        "        return spk_x2, mem_x2\n",
        "\n"
      ],
      "metadata": {
        "id": "U_D8NuZnxpXG"
      },
      "execution_count": 29,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# shape = (250, 1, 32, 32)\n",
        "\n",
        "# # Create a random tensor with the specified shape\n",
        "# x = torch.randn(shape)\n",
        "\n",
        "# print(x.size())\n",
        "# print(x.shape)\n",
        "\n",
        "# print(x.size(0))\n",
        "\n",
        "# print(x.view(x.size(0), -1).size())\n",
        "# x=x.view(x.size(0), -1)\n",
        "# print(\"x new.      \",x.size())\n",
        "\n",
        "# print(\"len(x)----------\",len(x))\n",
        "\n",
        "\n",
        "# # Assuming x is your tensor of size [250, 1024]\n",
        "# new_shape = (-1, 32)  # -1 means \"whatever is needed to keep the total number of elements the same\"\n",
        "# reshaped_x = x.view(new_shape)\n",
        "\n",
        "# print(\"x new.      \",x.size())\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "d3HCr9tn2JIi"
      },
      "execution_count": 30,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## ALL LEAKY"
      ],
      "metadata": {
        "id": "NQ7EsTgaHefc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# class SAE(nn.Module):\n",
        "#     def __init__(self):\n",
        "#         super().__init__()\n",
        "\n",
        "#              # Encoder\n",
        "#         self.enc0=    nn.Conv2d(1, 32, 3, padding=1, stride=2)\n",
        "#         self.enc1=    nn.BatchNorm2d(32)\n",
        "#         self.enc2=    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True, threshold=thresh)\n",
        "#         self.enc3=    nn.Conv2d(32, 64, 3, padding=1, stride=2)\n",
        "#         self.enc4=    nn.BatchNorm2d(64)\n",
        "#         self.enc5=    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True, threshold=thresh)\n",
        "#         self.enc6=    nn.Conv2d(64, 128, 3, padding=1, stride=2)\n",
        "#         self.enc7=    nn.BatchNorm2d(128)\n",
        "#         self.enc8=    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True, threshold=thresh)\n",
        "#         self.enc9=    nn.Flatten(start_dim=1, end_dim=3)\n",
        "#         self.enc10=   nn.Linear(128 * 4 * 4, latent_dim)\n",
        "#         self.enc11=   snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True, threshold=thresh)\n",
        "\n",
        "#         self.latent_dim = latent_dim\n",
        "\n",
        "#         self.linearNet0= nn.Linear(latent_dim,128*4*4)\n",
        "#         self.linearNet1=snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True,threshold=thresh)\n",
        "\n",
        "\n",
        "#         # Decoder:\n",
        "#         self.dec0=    nn.Unflatten(1,(128,4,4)) #Unflatten data from 1 dim to tensor of 128 x 4 x 4\n",
        "#         self.dec1=    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True,threshold=thresh)\n",
        "#         self.dec2=    nn.ConvTranspose2d(128, 64, 3,padding = 1,stride=(2,2),output_padding=1)\n",
        "#         self.dec3=    nn.BatchNorm2d(64)\n",
        "#         self.dec4=    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True, threshold=thresh)\n",
        "#         self.dec5=    nn.ConvTranspose2d(64, 32, 3,padding = 1,stride=(2,2),output_padding=1)\n",
        "#         self.dec6=    nn.BatchNorm2d(32)\n",
        "#         self.dec7=    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True, threshold=thresh)\n",
        "#         self.dec8=    nn.ConvTranspose2d(32, 1, 3,padding = 1,stride=(2,2),output_padding=1)\n",
        "#         self.dec9=    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,output=True,threshold=20000) #so membrane can be trained\n",
        "#        # self.dec9=    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,output=True,threshold=thresh) #---------------------------------------------- ADDED\n",
        "\n",
        "\n",
        "#     def forward(self, x):\n",
        "#         utils.reset(self.enc0)\n",
        "#         utils.reset(self.enc1)\n",
        "#         utils.reset(self.enc2)\n",
        "#         utils.reset(self.enc3)\n",
        "#         utils.reset(self.enc4)\n",
        "#         utils.reset(self.enc5)\n",
        "#         utils.reset(self.enc6)\n",
        "#         utils.reset(self.enc7)\n",
        "#         utils.reset(self.enc8)\n",
        "#         utils.reset(self.enc9)\n",
        "#         utils.reset(self.enc10)\n",
        "#         utils.reset(self.enc11)\n",
        "#         utils.reset(self.linearNet0)\n",
        "#         utils.reset(self.linearNet1)\n",
        "#         utils.reset(self.dec0)\n",
        "#         utils.reset(self.dec1)\n",
        "#         utils.reset(self.dec2)\n",
        "#         utils.reset(self.dec3)\n",
        "#         utils.reset(self.dec4)\n",
        "#         utils.reset(self.dec5)\n",
        "#         utils.reset(self.dec6)\n",
        "#         utils.reset(self.dec7)\n",
        "#         utils.reset(self.dec8)\n",
        "#         utils.reset(self.dec9)\n",
        "\n",
        "#         # # -------Initialize hidden states at t=0\n",
        "\n",
        "#         x2_rec = []\n",
        "#         x5_rec = []\n",
        "#         x8_rec = []\n",
        "#         x11_rec = []\n",
        "\n",
        "#         d1_rec = []\n",
        "#         d4_rec = []\n",
        "#         d7_rec = []\n",
        "#         d9_rec = []\n",
        "#         d9_rec_mem = []\n",
        "\n",
        "#         for step in range(num_steps):\n",
        "#             #------------------------------ encoder:\n",
        "#             x0=self.enc0(x)\n",
        "#             x1=self.enc1(x0)\n",
        "#             x2=self.enc2(x1) #Leaky\n",
        "#             # print(\"len x2: \", len(x2))\n",
        "#             # print(\"x2.size: \", x2.size())\n",
        "#             x3=self.enc3(x2)\n",
        "#             x4=self.enc4(x3)\n",
        "#             x5=self.enc5(x4) #Leaky\n",
        "#             x6=self.enc6(x5)\n",
        "#             x7=self.enc7(x6)\n",
        "#             x8=self.enc8(x7) #Leaky\n",
        "#             x9=self.enc9(x8)\n",
        "#             x10=self.enc10(x9)\n",
        "#             x11=self.enc11(x10) #Leaky\n",
        "#             # print(\"len x11: \", len(x11))\n",
        "#             # print(\"x11[0].size: \", x11[0].size())         torch.Size([250, 32])\n",
        "#             # print(\"x11[1].size: \", x11[1].size())         torch.Size([250, 32])\n",
        "\n",
        "\n",
        "#             x2_rec.append(x2)              #    x2 size:  torch.Size([250, 32, 16, 16])\n",
        "#             x5_rec.append(x5)                 #x5 size:  torch.Size([250, 64, 8, 8])\n",
        "#             x8_rec.append(x8)              #  x8 size:  torch.Size([250, 128, 4, 4])\n",
        "#             x11_rec.append(x11[0])\n",
        "#             # print(\"x11 length: \",len(x11))                        x11 length:  2\n",
        "#             # print(\"x11[1] length: \",len(x11[1]))               #   x11[1] length:  250\n",
        "#             # print(\"x11[0] length: \",len(x11[0]))               #   x11[0] length:  250\n",
        "#             # print(\"x11_rec length: \",len(x11_rec))                x11_rec length:  1 >> 2 >> 3 >> 4 >> 5\n",
        "\n",
        "#             #-------------------------------------------------\n",
        "#         x2_rec = torch.stack(x2_rec, dim=2)\n",
        "#         x5_rec = torch.stack(x5_rec, dim=2)\n",
        "#         x8_rec = torch.stack(x8_rec, dim=2)\n",
        "#         x11_rec = torch.stack(x11_rec, dim=2)                         #       torch.Size([250, 32, 5])\n",
        "\n",
        "#         x2_rec = x2_rec[:, :, -1]                                               # =========  x2_rec[:, :, -1].size() torch.Size([250, 32, 16, 16])\n",
        "#         x5_rec = x5_rec[:, :, -1]                                               # =========  x5_rec[:, :, -1].size() torch.Size([250, 64, 8, 8])\n",
        "#         # print(\"=========  x5_rec[:, :, -1].size()\", x5_rec.size())\n",
        "#         x8_rec = x8_rec[:, :, -1]                                               #  ========= x8_rec[:, :, -1].size() torch.Size([250, 128, 4, 4])\n",
        "#         # print(\"=========  x8_rec[:, :, -1].size()\", x8_rec.size())\n",
        "#         out_en = x11_rec[...,step]\n",
        "#         # print(\"========= out_en: x11_rec[...,step] size: \", out_en.size())\n",
        "\n",
        "#         # x11_rec = x11_rec[:, :, -1]             #  =========  x11_rec[:, :, -1].size() torch.Size([250, 32])\n",
        "\n",
        "#         for step in range(num_steps):\n",
        "#             #------------------------------ latent:\n",
        "#             # LN0=self.linearNet0(x11[...,step])         # means keep all dimensions of x11 up to the last one (which is usually the time dimension), and then selecting the data at the step position along that dimension.\n",
        "#             LN0=self.linearNet0(x11_rec[...,step])\n",
        "#             # print(\" LN0=self.linearNet0(x11_rec[...,step])======\", LN0.size())       # torch.Size([250, 2048])\n",
        "#             # print(\"type of LN0\",type(LN0))\n",
        "#             # print(\"length of LN0\",len(LN0))\n",
        "#             LN1=self.linearNet1(LN0)\n",
        "#             # print(\"LN1=self.linearNet1(LN0) size:\", len(LN1))\n",
        "#             # print(\"LN1[0].size()\", LN1[0].size())          #  LN1[0].size() torch.Size([250, 2048])\n",
        "#             # print(\"LN1[1].size()\", LN1[1].size())          #  LN1[1].size() torch.Size([250, 2048])\n",
        "\n",
        "#             # print(\"type of LN1\",type(LN1))\n",
        "#             # print(\"---------------length of LN1\",len(LN1))                      # length of LN1: 2\n",
        "#             # print(\"length of LN1\",len(LN1[0]))                                  # length of LN1[0]: 250\n",
        "#             # print(\"length of LN1\",len(LN1[1]))                                  # length of LN1[1]: 250\n",
        "#             #------------------------------ decoder:\n",
        "#             # d0=self.dec0(LN1[...,step])\n",
        "#             d0=self.dec0(LN1[0])\n",
        "#             # print(\"type of d0\",type(d0))\n",
        "#             # print(\"length of d0\",len(d0))\n",
        "#             # print(\"Size of d0\", LN1[0].size())              # Size of d0 torch.Size([250, 128, 4, 4])\n",
        "#             # print(\"Size of d0\", LN1[1].size())              # Size of d0 torch.Size([250, 128, 4, 4])\n",
        "\n",
        "#             d1=self.dec1(d0)     #Leaky\n",
        "#             d2=self.dec2(d1)\n",
        "#             d3=self.dec3(d2)\n",
        "#             d4=self.dec4(d3)     #Leaky\n",
        "#             d5=self.dec5(d4)\n",
        "#             d6=self.dec6(d5)\n",
        "#             d7=self.dec7(d6)     #Leaky\n",
        "#             d8=self.dec8(d7)\n",
        "#             d9=self.dec9(d8)     #Leaky\n",
        "#             # d9,d9_mem =self.dec9(d8)     #Leaky\n",
        "#             # print(\"            d9[0]=self.dec9(d8)\", d9[0].size())  #   torch.Size([250, 1, 32, 32])\n",
        "#             # print(\"            d9[1]=self.dec9(d8)\", d9[1].size())  #   torch.Size([250, 1, 32, 32])\n",
        "\n",
        "#             # x_recon=self.dec9(d8)\n",
        "#             #----------------------------------------\n",
        "\n",
        "#             #-----------decoder:\n",
        "#             d1_rec.append(d1)\n",
        "#             # print(\"length of d1_rec.append(d1)\",len(d1_rec))\n",
        "#             d4_rec.append(d4)\n",
        "#             # print(\"length of d4_rec.append(d4)\",len(d4_rec))\n",
        "#             d7_rec.append(d7)\n",
        "#             # print(\"length of d7_rec.append(d7)\",len(d7_rec))\n",
        "#             d9_rec.append(d9[0])\n",
        "#             # print(\"-----------d9 length: \", len(d9))                                   # length of d9: 2\n",
        "#             # print(\"d9[0] length: \", len(d9[0]))                                   # length of d9[0]: 250\n",
        "#             # print(\"d9[1] length: \", len(d9[1]))                                   # length of d9[1]: 250\n",
        "#             # print(\"length of d9_rec.append(d9)\",len(d9_rec))\n",
        "#             d9_rec_mem.append(d9[1])\n",
        "\n",
        "\n",
        "#        # decoder\n",
        "#         d1_rec = torch.stack(d1_rec, dim=4)                                              # torch.Size([250, 128, 4, 4, 5])\n",
        "#         # print(\"-----------------length of torch.stack(d1_rec, dim=4)\",len(d1_rec))\n",
        "#         # print(\"------size of torch.stack(d1_rec, dim=4)\",d1_rec.size())\n",
        "#         d4_rec = torch.stack(d4_rec, dim=4)                                              #torch.Size([250, 64, 8, 8, 5])\n",
        "#         # print(\"-----------------length of torch.stack(d4_rec, dim=4)\",len(d4_rec))\n",
        "#         # print(\"------size of torch.stack(d4_rec, dim=4)\",d4_rec.size())\n",
        "\n",
        "#         d7_rec = torch.stack(d7_rec, dim=4)                                              #torch.Size([250, 32, 16, 16, 5])\n",
        "#         # print(\"-----------------length of torch.stack(d7_rec, dim=4)\",len(d7_rec))\n",
        "#         # print(\"------size of torch.stack(d7_rec, dim=4)\",d7_rec.size())\n",
        "\n",
        "#         d9_rec = torch.stack(d9_rec, dim=4)\n",
        "#         # print(\"-----------------length of torch.stack(d9_rec, dim=4)\",len(d9_rec))\n",
        "#         # print(\"------size of torch.stack(d9_rec, dim=4)\",d9_rec.size())                  #size of torch.stack(d9_rec, dim=4) torch.Size([250, 1, 32, 32, 5])\n",
        "#         d9_rec_mem = torch.stack(d9_rec_mem, dim=4)\n",
        "\n",
        "#        # decoder\n",
        "#         d1_rec = d1_rec[:, :, :, :, -1]\n",
        "#         # print(\"=========  d1_rec[:, :, :, :, -1].size()\", d1_rec.size())              #   d1_rec[:, :, :, :, -1].size() torch.Size([250, 128, 4, 4])\n",
        "#         d4_rec = d4_rec[:, :, :, :, -1]\n",
        "#         # print(\"=========  d4_rec[:, :, :, :, -1].size()\", d4_rec.size())              #   d4_rec[:, :, :, :, -1].size() torch.Size([250, 64, 8, 8])\n",
        "#         d7_rec = d7_rec[:, :, :, :, -1]\n",
        "#         # print(\"=========  d7_rec[:, :, :, :, -1].size()\", d7_rec.size())              #   d7_rec[:, :, :, :, -1].size() torch.Size([250, 32, 16, 16])\n",
        "#         d9_rec = d9_rec[:, :, :, :, -1]\n",
        "#         # print(\"=========  d9_rec[:, :, :, :, -1].size()\", d9_rec.size())              #   d9_rec[:, :, :, :, -1].size() torch.Size([250, 1, 32, 32])\n",
        "#         out = d9_rec_mem[:, :, :, :, -1]\n",
        "\n",
        "#         # self.out_en = out_en\n",
        "#         # self.out = out\n",
        "\n",
        "#         return d9_rec, out_en\n",
        "\n",
        "#     def get_activation(self, name):\n",
        "#         def hook(module, input, output):\n",
        "#             setattr(self, name, output)  # Store the output as an attribute of the model\n",
        "#         return hook\n",
        "\n",
        "\n",
        "# # activation = {}\n",
        "# #     def get_activation(self, name):\n",
        "# #         def hook(model, input, output):\n",
        "# #             activation[name] = output.detach()\n",
        "# #         return hook\n"
      ],
      "metadata": {
        "id": "AIVR3KS_XFcu"
      },
      "execution_count": 31,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# class SAE(nn.Module):\n",
        "#     def __init__(self):\n",
        "#         super().__init__()\n",
        "\n",
        "#              # Encoder\n",
        "#         self.enc0=    nn.Conv2d(1, 32, 3, padding=1, stride=2)\n",
        "#         self.enc1=    nn.BatchNorm2d(32)\n",
        "#         self.enc2=    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, threshold=thresh)\n",
        "#         self.enc3=    nn.Conv2d(32, 64, 3, padding=1, stride=2)\n",
        "#         self.enc4=    nn.BatchNorm2d(64)\n",
        "#         self.enc5=    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True , threshold=thresh)\n",
        "#         self.enc6=    nn.Conv2d(64, 128, 3, padding=1, stride=2)\n",
        "#         self.enc7=    nn.BatchNorm2d(128)\n",
        "#         self.enc8=    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, threshold=thresh)\n",
        "#         self.enc9=    nn.Flatten(start_dim=1, end_dim=3)\n",
        "#         self.enc10=   nn.Linear(128 * 4 * 4, latent_dim)\n",
        "#         self.enc11=   snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True, threshold=thresh)\n",
        "\n",
        "#         self.latent_dim = latent_dim\n",
        "\n",
        "#         self.linearNet0= nn.Linear(latent_dim,128*4*4)\n",
        "#         self.linearNet1=snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True,threshold=thresh)\n",
        "\n",
        "\n",
        "#         # Decoder:\n",
        "#         self.dec0=    nn.Unflatten(1,(128,4,4)) #Unflatten data from 1 dim to tensor of 128 x 4 x 4\n",
        "#         self.dec1=    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,threshold=thresh)\n",
        "#         self.dec2=    nn.ConvTranspose2d(128, 64, 3,padding = 1,stride=(2,2),output_padding=1)\n",
        "#         self.dec3=    nn.BatchNorm2d(64)\n",
        "#         self.dec4=    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, threshold=thresh)\n",
        "#         self.dec5=    nn.ConvTranspose2d(64, 32, 3,padding = 1,stride=(2,2),output_padding=1)\n",
        "#         self.dec6=    nn.BatchNorm2d(32)\n",
        "#         self.dec7=    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, threshold=thresh)\n",
        "#         self.dec8=    nn.ConvTranspose2d(32, 1, 3,padding = 1,stride=(2,2),output_padding=1)\n",
        "#         self.dec9=    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,output=True,threshold=20000) #so membrane can be trained\n",
        "#        # self.dec9=    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True,output=True,threshold=thresh) #---------------------------------------------- ADDED\n",
        "\n",
        "\n",
        "#     def forward(self, x):\n",
        "#         utils.reset(self.enc0)\n",
        "#         utils.reset(self.enc1)\n",
        "#         utils.reset(self.enc2)\n",
        "#         utils.reset(self.enc3)\n",
        "#         utils.reset(self.enc4)\n",
        "#         utils.reset(self.enc5)\n",
        "#         utils.reset(self.enc6)\n",
        "#         utils.reset(self.enc7)\n",
        "#         utils.reset(self.enc8)\n",
        "#         utils.reset(self.enc9)\n",
        "#         utils.reset(self.enc10)\n",
        "#         utils.reset(self.enc11)\n",
        "#         utils.reset(self.linearNet0)\n",
        "#         utils.reset(self.linearNet1)\n",
        "#         utils.reset(self.dec0)\n",
        "#         utils.reset(self.dec1)\n",
        "#         utils.reset(self.dec2)\n",
        "#         utils.reset(self.dec3)\n",
        "#         utils.reset(self.dec4)\n",
        "#         utils.reset(self.dec5)\n",
        "#         utils.reset(self.dec6)\n",
        "#         utils.reset(self.dec7)\n",
        "#         utils.reset(self.dec8)\n",
        "#         utils.reset(self.dec9)\n",
        "\n",
        "\n",
        "#         x2_rec = []\n",
        "#         x5_rec = []\n",
        "#         x8_rec = []\n",
        "#         x11_rec = []\n",
        "\n",
        "#         d1_rec = []\n",
        "#         d4_rec = []\n",
        "#         d7_rec = []\n",
        "#         d9_rec = []\n",
        "#         d9_rec_mem = []\n",
        "\n",
        "#         for step in range(num_steps):\n",
        "#             #------------------------------ encoder:\n",
        "#             x0=self.enc0(x)\n",
        "#             x1=self.enc1(x0)\n",
        "#             x2=self.enc2(x1) #Leaky\n",
        "#             x3=self.enc3(x2)\n",
        "#             x4=self.enc4(x3)\n",
        "#             x5, x5mem =self.enc5(x4) #Leaky\n",
        "#             x6=self.enc6(x5)\n",
        "#             x7=self.enc7(x6)\n",
        "#             x8=self.enc8(x7) #Leaky\n",
        "#             x9=self.enc9(x8)\n",
        "#             x10=self.enc10(x9)\n",
        "#             x11=self.enc11(x10) #Leaky\n",
        "\n",
        "#             x2_rec.append(x2)\n",
        "#             x5_rec.append(x5)\n",
        "#             x8_rec.append(x8)\n",
        "#             x11_rec.append(x11[0])\n",
        "\n",
        "#             #-------------------------------------------------\n",
        "#         x2_rec = torch.stack(x2_rec, dim=2)\n",
        "#         x5_rec = torch.stack(x5_rec, dim=2)\n",
        "#         x8_rec = torch.stack(x8_rec, dim=2)\n",
        "#         x11_rec = torch.stack(x11_rec, dim=2)                         #\n",
        "#         x2_rec = x2_rec[:, :, -1]\n",
        "#         x5_rec = x5_rec[:, :, -1]\n",
        "#         x8_rec = x8_rec[:, :, -1]\n",
        "#         out_en = x11_rec[...,step]\n",
        "#         # x11_rec = x11_rec[:, :, -1]             #  =========  x11_rec[:, :, -1].size() torch.Size([250, 32])\n",
        "\n",
        "#         for step in range(num_steps):\n",
        "#             #------------------------------ latent:\n",
        "#             LN0=self.linearNet0(x11_rec[...,step])\n",
        "#             LN1=self.linearNet1(LN0)\n",
        "#             #------------------------------ decoder:\n",
        "#             # d0=self.dec0(LN1[...,step])\n",
        "#             d0=self.dec0(LN1[0])\n",
        "#             d1=self.dec1(d0)     #Leaky\n",
        "#             d2=self.dec2(d1)\n",
        "#             d3=self.dec3(d2)\n",
        "#             d4=self.dec4(d3)     #Leaky\n",
        "#             d5=self.dec5(d4)\n",
        "#             d6=self.dec6(d5)\n",
        "#             d7=self.dec7(d6)     #Leaky\n",
        "#             d8=self.dec8(d7)\n",
        "#             d9=self.dec9(d8)     #Leaky\n",
        "#             #----------------------------------------\n",
        "\n",
        "#             #-----------decoder:\n",
        "#             d1_rec.append(d1)\n",
        "#             d4_rec.append(d4)\n",
        "#             d7_rec.append(d7)\n",
        "#             d9_rec.append(d9[0])\n",
        "#             d9_rec_mem.append(d9[1])\n",
        "\n",
        "\n",
        "#        # decoder\n",
        "#         d1_rec = torch.stack(d1_rec, dim=4)\n",
        "#         d4_rec = torch.stack(d4_rec, dim=4)\n",
        "#         d7_rec = torch.stack(d7_rec, dim=4)\n",
        "#         d9_rec = torch.stack(d9_rec, dim=4)\n",
        "#         d9_rec_mem = torch.stack(d9_rec_mem, dim=4)\n",
        "\n",
        "#        # decoder\n",
        "#         d1_rec = d1_rec[:, :, :, :, -1]\n",
        "#         d4_rec = d4_rec[:, :, :, :, -1]\n",
        "#         d7_rec = d7_rec[:, :, :, :, -1]\n",
        "#         d9_rec = d9_rec[:, :, :, :, -1]\n",
        "#         out = d9_rec_mem[:, :, :, :, -1]\n",
        "\n",
        "#         # self.out_en = out_en\n",
        "#         # self.out = out\n",
        "\n",
        "#         return d9_rec, out_en\n",
        "\n",
        "#     # def get_activation(self, name):\n",
        "#     #     def hook(module, input, output):\n",
        "#     #         setattr(self, name, output)  # Store the output as an attribute of the model\n",
        "#     #     return hook"
      ],
      "metadata": {
        "id": "_C76N8tQJmcT"
      },
      "execution_count": 32,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## for each layer"
      ],
      "metadata": {
        "id": "9q4mUUpWdF_a"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# # Training and Testing\n",
        "# from torchvision.utils import save_image\n",
        "\n",
        "# spike_recordings = []\n",
        "# train_ber_rec = []\n",
        "# test_ber_rec = []\n",
        "# threshold_Real = 0.5\n",
        "# threshold_Recon = 0.5\n",
        "\n",
        "# def train(network, trainloader, opti, epoch):\n",
        "#     network=network.train()\n",
        "#     train_loss_hist=torch.zeros((1), dtype=dtype, device=device)\n",
        "#     train_avg_loss_rec=[]\n",
        "\n",
        "#     for batch_idx, (real_img, labels) in enumerate(trainloader):\n",
        "#         opti.zero_grad()\n",
        "#         real_img = real_img.to(device)\n",
        "#         labels = labels.to(device)\n",
        "\n",
        "#         # print(\"real_img size\", real_img.size())      #    real_img size                # -------------------------------------------------ADDED\n",
        "#         out, out_en = network(real_img)\n",
        "#         x_recon, out_en = network(real_img)\n",
        "#         # print(\"out_en size\", out_en.size())      #                  # -------------------------------------------------ADDED\n",
        "#         # print(\"x_recon size\", x_recon.size())      #                  # -------------------------------------------------ADDED\n",
        "\n",
        "\n",
        "#         #Calculate loss\n",
        "#         loss_val = torch.zeros((1), dtype=dtype, device=device)\n",
        "#         for step in range(num_steps):\n",
        "#           loss_val += F.mse_loss(x_recon, real_img)                  #.view(1, -1)\n",
        "\n",
        "#         train_loss_hist += (loss_val.item())/num_steps\n",
        "#         avg_loss=train_loss_hist.mean()\n",
        "\n",
        "#         # # # ---------------------\n",
        "#         print(f'Train[{epoch}/{max_epoch}][{batch_idx}/{len(trainloader)}] Loss: {loss_val.item()}')\n",
        "\n",
        "#         loss_val.backward()        #\n",
        "\n",
        "#         opti.step()\n",
        "#         train_loss_rec.append(loss_val.item())\n",
        "\n",
        "#         #Save reconstructed images every at the end of the epoch\n",
        "#         if batch_idx == len(trainloader)-1:\n",
        "#             utls.save_image((real_img+1)/2, f'figures/training/epoch{epoch}_finalbatch_inputs.png')\n",
        "#             utls.save_image((x_recon+1)/2, f'figures/training/epoch{epoch}_finalbatch_recon.png')\n",
        "#             train_auc = auc(np.arange(len(train_loss_rec)), train_loss_rec)\n",
        "\n",
        "#     # return loss_val, train_loss_rec, train_auc, d9_rec, out_en  #              # -------------------------------------------------ADDED\n",
        "#     return loss_val, train_loss_rec, train_auc, out, out_en  #              # -------------------------------------------------ADDED\n",
        "\n",
        "\n",
        "# # For Testing, not doing backpropagate, therefore no gradients are required and we use torch.no_grad():\n",
        "# #Testing Loop\n",
        "# def test(network, testloader, opti, epoch):\n",
        "#     network=network.eval()\n",
        "#     test_loss_hist=[]\n",
        "#     test_avg_loss_rec=[]\n",
        "#     test_avg_loss_hist = []\n",
        "\n",
        "#     spk_rec_test = []\n",
        "#     with torch.no_grad(): #no gradient this time\n",
        "#         for batch_idx, (real_img, labels) in enumerate(testloader):\n",
        "#             real_img = real_img.to(device)#\n",
        "#             labels = labels.to(device)\n",
        "#             # x11_rec, d9_rec = network(real_img)             # -------------------------------------------------ADDED\n",
        "#             out, out_en = network(real_img)  # Pass data into network and return reconstructed image and spk_rec\n",
        "#             x_recon, out_en = network(real_img)  # Pass data into network and return reconstructed image and spk_rec\n",
        "#             # average Loss:\n",
        "#             loss_val = torch.zeros((1), dtype=dtype, device=device)\n",
        "#             for step in range(num_steps):\n",
        "#               loss_val += F.mse_loss(x_recon, real_img)\n",
        "#             avg_loss=loss_val/num_steps\n",
        "#             test_loss_hist.append(loss_val.item())\n",
        "\n",
        "#             real_img_binary = (real_img > threshold_Real).float()\n",
        "#             x_recon_binary = (x_recon > threshold_Recon).float()\n",
        "#             bit_errors = torch.sum(torch.abs(real_img_binary - x_recon_binary))\n",
        "#             total_pixels = real_img_binary.numel()  # Total number of pixels in the images\n",
        "#             bit_error_rate = bit_errors.item() / total_pixels\n",
        "#             test_ber_rec.append(bit_error_rate)  # Append BER to the list\n",
        "\n",
        "#             # Save binary images\n",
        "#             Error_bin = (torch.abs(x_recon_binary - real_img_binary))\n",
        "\n",
        "#             if batch_idx == len(testloader)-1:\n",
        "#               if epoch in [0, 25, 49]:\n",
        "#                 save_image(real_img_binary, f'figures/binarytesting/ep{epoch}_inputs_binary.png')\n",
        "#                 save_image(x_recon_binary, f'figures/binarytesting/ep{epoch}_recon_binary.png')\n",
        "#                 save_image(Error_bin, f'figures/binarytesting/ep{epoch}_Error_bin.png')\n",
        "#             # -------------------------------------------------------------------------------------------------\n",
        "\n",
        "#             print(f'Test[{epoch}/{max_epoch}][{batch_idx}/{len(testloader)}]  Loss: {loss_val.item()}, '  f'BER (test): {bit_error_rate}')\n",
        "\n",
        "#             test_loss_rec.append(loss_val.item())\n",
        "\n",
        "#             if batch_idx == len(testloader)-1:\n",
        "#                 utls.save_image((real_img+1)/2, f'figures/testing/epoch{epoch}_finalbatch_inputs.png')\n",
        "#                 utls.save_image((x_recon+1)/2, f'figures/testing/epoch{epoch}_finalbatch_recons.png')\n",
        "#                 test_auc = auc(np.arange(len(test_loss_rec)), test_loss_rec)\n",
        "\n",
        "#     # return loss_val, test_loss_rec, test_auc, out, out_en                     # -------------------------------------------------ADDED\n",
        "#     # return loss_val, test_loss_rec, test_auc, x11_rec, d9_rec  #              # -------------------------------------------------ADDED\n",
        "#     return loss_val, test_loss_rec, test_auc, out, out_en\n",
        "\n"
      ],
      "metadata": {
        "id": "PGWPHsrmPW9v"
      },
      "execution_count": 33,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# # Training and Testing\n",
        "# # using MSE loss to compare the reconstructed image (x_recon) with the original image (real_img)\n",
        "# from torchvision.utils import save_image\n",
        "\n",
        "# spike_recordings = []\n",
        "# train_ber_rec = []\n",
        "# test_ber_rec = []\n",
        "# threshold_Real = 0.5\n",
        "# threshold_Recon = 0.5\n",
        "\n",
        "# def train(network, trainloader, opti, epoch):\n",
        "#     network=network.train()\n",
        "#     train_loss_hist=torch.zeros((1), dtype=dtype, device=device)\n",
        "#     train_avg_loss_rec=[]\n",
        "\n",
        "#     for batch_idx, (real_img, labels) in enumerate(trainloader):\n",
        "#         opti.zero_grad()\n",
        "#         real_img = real_img.to(device)\n",
        "#         labels = labels.to(device)\n",
        "#         # d9_rec, out_en = network(real_img)             # -------------------------------------------------ADDED\n",
        "#         # print(\"real_img size\", real_img.size())      #    real_img size torch.Size([250, 1, 32, 32])                # -------------------------------------------------ADDED\n",
        "\n",
        "#         # out, out_en = network(real_img)  # Pass data into network and return reconstructed image and spk_rec\n",
        "#         x_recon, out_en = network(real_img)  # Pass data into network and return reconstructed image and spk_rec\n",
        "#         # print(\"x_recon size\", x_recon.size())      #       x_recon size torch.Size([250, 1, 32, 32])\n",
        "#         # print(\"out size\", out.size())             #       out size torch.Size([250, 32])\n",
        "\n",
        "#         #Calculate loss\n",
        "#         loss_val = torch.zeros((1), dtype=dtype, device=device)\n",
        "#         for step in range(num_steps):\n",
        "#           loss_val += F.mse_loss(x_recon, real_img)                  #.view(1, -1)\n",
        "\n",
        "#             # Clone the loss_val tensor to avoid in-place modification\n",
        "\n",
        "#         train_loss_hist += (loss_val.item())/num_steps\n",
        "#         avg_loss=train_loss_hist.mean()\n",
        "\n",
        "#         # # ---------------------------- Calculate Bit Error Rate (BER)\n",
        "#         real_img_binary = (real_img > threshold_Real).float()\n",
        "#         x_recon_binary = (x_recon > threshold_Recon).float()\n",
        "#         bit_errors = torch.sum(torch.abs(real_img_binary - x_recon_binary))\n",
        "#         total_pixels = real_img_binary.numel()  # Total number of pixels in the images\n",
        "#         bit_error_rate = bit_errors.item() / total_pixels\n",
        "#         train_ber_rec.append(bit_error_rate)  # Append BER to the list\n",
        "\n",
        "#         # Save binary images\n",
        "#         Error_bin = (torch.abs(x_recon_binary - real_img_binary))\n",
        "#         if batch_idx == len(trainloader)-1:\n",
        "#           if epoch in [0, 25, 49]:\n",
        "#             utls.save_image(real_img_binary, f'figures/binarytraining/ep{epoch}_inputs_binary.png')\n",
        "#             utls.save_image(x_recon_binary, f'figures/binarytraining/ep{epoch}_recon_binary.png')\n",
        "#             utls.save_image(Error_bin, f'figures/binarytraining/ep{epoch}_Error_bin.png')\n",
        "#         print(f'Train[{epoch}/{max_epoch}][{batch_idx}/{len(trainloader)}] Loss: {loss_val.item()}, ' f'BER : {bit_error_rate}')\n",
        "\n",
        "#         # loss_val += torch.mean(loss_val)  # Accumulate the loss             # -------------------------------------------------ADDED\n",
        "#         loss_val.backward()         #\n",
        "#         # loss_val.backward(retain_graph=True)        #\n",
        "\n",
        "#         opti.step()\n",
        "#         train_loss_rec.append(loss_val.item())\n",
        "\n",
        "#         #Save reconstructed images every at the end of the epoch\n",
        "#         if batch_idx == len(trainloader)-1:\n",
        "#             utls.save_image((real_img+1)/2, f'figures/training/epoch{epoch}_finalbatch_inputs.png')\n",
        "#             utls.save_image((x_recon+1)/2, f'figures/training/epoch{epoch}_finalbatch_recon.png')\n",
        "#             train_auc = auc(np.arange(len(train_loss_rec)), train_loss_rec)\n",
        "\n",
        "#     # loss_val.backward()                                            # -------------------------------------------------ADDED\n",
        "\n",
        "#     # return loss_val, train_loss_rec, train_auc , out, out_en  #, spk_rec_batches#, train_avg_loss_rec, #avg_loss #, train_loss_hist\n",
        "#     # return loss_val, train_loss_rec, train_auc   #\n",
        "#     return loss_val, train_loss_rec, train_auc, d9_rec, out_en  #              # -------------------------------------------------ADDED\n",
        "\n",
        "\n",
        "# # For Testing, not doing backpropagate, therefore no gradients are required and we use torch.no_grad():\n",
        "# #Testing Loop\n",
        "# def test(network, testloader, opti, epoch):\n",
        "#     network=network.eval()\n",
        "#     test_loss_hist=[]\n",
        "#     test_avg_loss_rec=[]\n",
        "#     test_avg_loss_hist = []\n",
        "\n",
        "#     spk_rec_test = []\n",
        "#     with torch.no_grad(): #no gradient this time\n",
        "#         for batch_idx, (real_img, labels) in enumerate(testloader):\n",
        "#             real_img = real_img.to(device)#\n",
        "#             labels = labels.to(device)\n",
        "#             # x11_rec, d9_rec = network(real_img)             # -------------------------------------------------ADDED\n",
        "#             # out, out_en = network(real_img)  # Pass data into network and return reconstructed image and spk_rec\n",
        "#             x_recon, out_en = network(real_img)  # Pass data into network and return reconstructed image and spk_rec\n",
        "#             # average Loss:\n",
        "#             loss_val = torch.zeros((1), dtype=dtype, device=device)\n",
        "#             for step in range(num_steps):\n",
        "#               loss_val += F.mse_loss(x_recon, real_img)\n",
        "#             avg_loss=loss_val/num_steps\n",
        "#             test_loss_hist.append(loss_val.item())\n",
        "\n",
        "#             real_img_binary = (real_img > threshold_Real).float()\n",
        "#             x_recon_binary = (x_recon > threshold_Recon).float()\n",
        "#             bit_errors = torch.sum(torch.abs(real_img_binary - x_recon_binary))\n",
        "#             total_pixels = real_img_binary.numel()  # Total number of pixels in the images\n",
        "#             bit_error_rate = bit_errors.item() / total_pixels\n",
        "#             test_ber_rec.append(bit_error_rate)  # Append BER to the list\n",
        "\n",
        "#             # Save binary images\n",
        "#             Error_bin = (torch.abs(x_recon_binary - real_img_binary))\n",
        "\n",
        "#             if batch_idx == len(testloader)-1:\n",
        "#               if epoch in [0, 25, 49]:\n",
        "#                 save_image(real_img_binary, f'figures/binarytesting/ep{epoch}_inputs_binary.png')\n",
        "#                 save_image(x_recon_binary, f'figures/binarytesting/ep{epoch}_recon_binary.png')\n",
        "#                 save_image(Error_bin, f'figures/binarytesting/ep{epoch}_Error_bin.png')\n",
        "#             # -------------------------------------------------------------------------------------------------\n",
        "\n",
        "#             print(f'Test[{epoch}/{max_epoch}][{batch_idx}/{len(testloader)}]  Loss: {loss_val.item()}, '  f'BER (test): {bit_error_rate}')\n",
        "\n",
        "#             test_loss_rec.append(loss_val.item())\n",
        "\n",
        "#             if batch_idx == len(testloader)-1:\n",
        "#                 utls.save_image((real_img+1)/2, f'figures/testing/epoch{epoch}_finalbatch_inputs.png')\n",
        "#                 utls.save_image((x_recon+1)/2, f'figures/testing/epoch{epoch}_finalbatch_recons.png')\n",
        "#                 test_auc = auc(np.arange(len(test_loss_rec)), test_loss_rec)\n",
        "\n",
        "#     # return loss_val, test_loss_rec, test_auc, out, out_en                     # -------------------------------------------------ADDED\n",
        "#     # return loss_val, test_loss_rec, test_auc, x11_rec, d9_rec  #              # -------------------------------------------------ADDED\n",
        "#     return loss_val, test_loss_rec, test_auc  #\n",
        "\n"
      ],
      "metadata": {
        "id": "HmvBPb95dEju"
      },
      "execution_count": 34,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "metadata": {
        "id": "5HbcSNNp6vK2"
      },
      "outputs": [],
      "source": [
        "# Training and Testing\n",
        "# using MSE loss to compare the reconstructed image (x_recon) with the original image (real_img)\n",
        "from torchvision.utils import save_image\n",
        "\n",
        "spike_recordings = []\n",
        "train_ber_rec = []\n",
        "test_ber_rec = []\n",
        "threshold_Real = 0.5\n",
        "threshold_Recon = 0.5\n",
        "std_dev=0.1\n",
        "\n",
        "\n",
        "def train(network, trainloader, opti, epoch):\n",
        "    network=network.train()\n",
        "    train_loss_hist=torch.zeros((1), dtype=dtype, device=device)\n",
        "    train_avg_loss_rec=[]\n",
        "\n",
        "    for batch_idx, (real_img, labels) in enumerate(trainloader):\n",
        "        opti.zero_grad()\n",
        "        real_img = real_img.to(device)\n",
        "        labels = labels.to(device)\n",
        "\n",
        "        # -------------Add Gaussian noise to the input images\n",
        "        real_img = real_img + torch.randn_like(real_img) * std_dev\n",
        "\n",
        "        out, out_en = network(real_img)  # Pass data into network and return reconstructed image and spk_rec\n",
        "        x_recon, out = network(real_img)  # Pass data into network and return reconstructed image and spk_rec.  #        x_recon size torch.Size([250, 1, 32, 32]) ,  #        out size torch.Size([250, 32])\n",
        "        #Calculate loss\n",
        "        loss_val = torch.zeros((1), dtype=dtype, device=device)\n",
        "        for step in range(num_steps):\n",
        "          loss_val += F.mse_loss(x_recon, real_img)                  #.view(1, -1)\n",
        "        train_loss_hist += (loss_val.item())/num_steps\n",
        "        avg_loss=train_loss_hist.mean()\n",
        "\n",
        "        # # ---------------------------- Calculate Bit Error Rate (BER)\n",
        "        real_img_binary = (real_img > threshold_Real).float()\n",
        "        x_recon_binary = (x_recon > threshold_Recon).float()\n",
        "        bit_errors = torch.sum(torch.abs(real_img_binary - x_recon_binary))\n",
        "        total_pixels = real_img_binary.numel()  # Total number of pixels in the images\n",
        "        bit_error_rate = bit_errors.item() / total_pixels\n",
        "        train_ber_rec.append(bit_error_rate)  # Append BER to the list\n",
        "\n",
        "        # Save binary images\n",
        "        Error_bin = (torch.abs(x_recon_binary - real_img_binary))\n",
        "\n",
        "        if batch_idx == len(trainloader)-1:\n",
        "          if epoch in [0, 25, 49]:\n",
        "            utls.save_image(real_img_binary, f'figures/binarytraining/ep{epoch}_inputs_binary.png')\n",
        "            utls.save_image(x_recon_binary, f'figures/binarytraining/ep{epoch}_recon_binary.png')\n",
        "            utls.save_image(Error_bin, f'figures/binarytraining/ep{epoch}_Error_bin.png')\n",
        "        print(f'Train[{epoch}/{max_epoch}][{batch_idx}/{len(trainloader)}] Loss: {loss_val.item()}, ' f'BER : {bit_error_rate}')\n",
        "\n",
        "        loss_val.backward()\n",
        "        opti.step()\n",
        "        train_loss_rec.append(loss_val.item())\n",
        "\n",
        "        #Save reconstructed images every at the end of the epoch\n",
        "        if batch_idx == len(trainloader)-1:\n",
        "            utls.save_image((real_img+1)/2, f'figures/training/epoch{epoch}_finalbatch_inputs.png')\n",
        "            utls.save_image((x_recon+1)/2, f'figures/training/epoch{epoch}_finalbatch_recon.png')\n",
        "            train_auc = auc(np.arange(len(train_loss_rec)), train_loss_rec)\n",
        "\n",
        "    return loss_val, train_loss_rec, train_auc , out, out_en  #, spk_rec_batches#, train_avg_loss_rec, #avg_loss #, train_loss_hist\n",
        "    # return loss_val, train_loss_rec, train_auc   #\n",
        "\n",
        "\n",
        "# For Testing, not doing backpropagate, therefore no gradients are required and we use torch.no_grad():\n",
        "#Testing Loop\n",
        "def test(network, testloader, opti, epoch):\n",
        "    network=network.eval()\n",
        "    test_loss_hist=[]\n",
        "    test_avg_loss_rec=[]\n",
        "    test_avg_loss_hist = []\n",
        "\n",
        "    spk_rec_test = []\n",
        "    with torch.no_grad(): #no gradient this time\n",
        "        for batch_idx, (real_img, labels) in enumerate(testloader):\n",
        "            real_img = real_img.to(device)#\n",
        "            labels = labels.to(device)\n",
        "            out, out_en = network(real_img)  # Pass data into network and return reconstructed image and spk_rec\n",
        "            x_recon , out = network(real_img)  # Pass data into network and return reconstructed image and spk_rec\n",
        "            # average Loss:\n",
        "            loss_val = torch.zeros((1), dtype=dtype, device=device)\n",
        "            for step in range(num_steps):\n",
        "              loss_val += F.mse_loss(x_recon, real_img)\n",
        "            avg_loss=loss_val/num_steps\n",
        "            test_loss_hist.append(loss_val.item())\n",
        "\n",
        "            real_img_binary = (real_img > threshold_Real).float()\n",
        "            x_recon_binary = (x_recon > threshold_Recon).float()\n",
        "            bit_errors = torch.sum(torch.abs(real_img_binary - x_recon_binary))\n",
        "            total_pixels = real_img_binary.numel()  # Total number of pixels in the images\n",
        "            bit_error_rate = bit_errors.item() / total_pixels\n",
        "            test_ber_rec.append(bit_error_rate)  # Append BER to the list\n",
        "\n",
        "            # Save binary images\n",
        "            Error_bin = (torch.abs(x_recon_binary - real_img_binary))\n",
        "\n",
        "            if batch_idx == len(testloader)-1:\n",
        "              if epoch in [0, 25, 49]:\n",
        "                save_image(real_img_binary, f'figures/binarytesting/ep{epoch}_inputs_binary.png')\n",
        "                save_image(x_recon_binary, f'figures/binarytesting/ep{epoch}_recon_binary.png')\n",
        "                save_image(Error_bin, f'figures/binarytesting/ep{epoch}_Error_bin.png')\n",
        "            # -------------------------------------------------------------------------------------------------\n",
        "\n",
        "            print(f'Test[{epoch}/{max_epoch}][{batch_idx}/{len(testloader)}]  Loss: {loss_val.item()}, '  f'BER (test): {bit_error_rate}')\n",
        "\n",
        "            test_loss_rec.append(loss_val.item())\n",
        "\n",
        "            if batch_idx == len(testloader)-1:\n",
        "                utls.save_image((real_img+1)/2, f'figures/testing/epoch{epoch}_finalbatch_inputs.png')\n",
        "                utls.save_image((real_img+(torch.randn_like(real_img) * std_dev)+1)/2, f'figures/testing/(real_img+noise)_epoch{epoch}_finalbatch_inputs.png')\n",
        "                utls.save_image((x_recon+1)/2, f'figures/testing/epoch{epoch}_finalbatch_recons.png')\n",
        "                test_auc = auc(np.arange(len(test_loss_rec)), test_loss_rec)\n",
        "\n",
        "    return loss_val, test_loss_rec, test_auc, out, out_en                     # -------------------------------------------------ADDED\n",
        "    # return loss_val, test_loss_rec, test_auc  #\n",
        "\n",
        "for batch_spikes in spike_recordings:\n",
        "    print(batch_spikes)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "metadata": {
        "id": "lF2cwM20PKAL"
      },
      "outputs": [],
      "source": [
        "input_size = 32 #resize of mnist data (optional)\n",
        "\n",
        "#setup GPU\n",
        "dtype = torch.float\n",
        "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
        "\n",
        "# neuron and simulation parameters\n",
        "spike_grad = surrogate.atan(alpha=2.0)  # alternate surrogate gradient fast_sigmoid(slope=25)\n",
        "\n",
        "train_loss_rec = []\n",
        "test_loss_rec = []\n",
        "train_loss_record = []\n",
        "test_loss_record = []\n",
        "train_avg_loss_rec=[]\n",
        "test_avg_loss_rec=[]\n",
        "\n",
        "  # Synaptic current and membrane potential decay exponentially with rates of alpha and beta\n",
        "alpha=0.9\n",
        "beta_syn=0.0001\n",
        "# beta_syn=0.9\n",
        "\n",
        "beta =0.9\n",
        "\n",
        "num_steps=5\n",
        "latent_dim = 32 #dimension of latent layer (how compressed we want the information)\n",
        "thresh=1    #spiking threshold (lower = more spikes are let through)\n",
        "epochs=50\n",
        "# epochs=5\n",
        "max_epoch=epochs\n",
        "\n",
        "  #Define Network and optimizer\n",
        "net=SAE()\n",
        "net = net.to(device)\n",
        "optimizer = torch.optim.AdamW(net.parameters(),\n",
        "                            lr=0.0001,\n",
        "                            betas=(0.9, 0.999),\n",
        "                            weight_decay=0.001)\n",
        "\n",
        "\n",
        "\n",
        "activation = {}\n",
        "# def get_activation(name):\n",
        "#     def hook(model, input, output):\n",
        "#         activation[name] = output.detach()\n",
        "#     return hook\n",
        "\n",
        "\n",
        "def get_activation(name):\n",
        "    def hook(model, input, output):\n",
        "        if isinstance(output, tuple):\n",
        "            activation[name] = [out.detach() for out in output]\n",
        "        else:\n",
        "            activation[name] = output.detach()\n",
        "    return hook\n",
        "\n",
        "\n",
        "# net.encoder[5].register_forward_hook(get_activation('encoder[5]'))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ja2hQ0wiOxUr"
      },
      "source": [
        "## for saving the out_en after each epoch:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cUOjPTrBOjmt",
        "outputId": "26d25690-c9ca-4524-bf58-8a8da39522c3"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Train[0/50][0/240] Loss: 8.006251335144043, BER : 0.86759765625\n",
            "Train[0/50][1/240] Loss: 8.01241683959961, BER : 0.86846875\n",
            "Train[0/50][2/240] Loss: 8.04263687133789, BER : 0.87267578125\n",
            "Train[0/50][3/240] Loss: 8.040471076965332, BER : 0.87218359375\n",
            "Train[0/50][4/240] Loss: 7.965984344482422, BER : 0.86468359375\n",
            "Train[0/50][5/240] Loss: 8.055792808532715, BER : 0.8756953125\n",
            "Train[0/50][6/240] Loss: 8.02262020111084, BER : 0.87193359375\n",
            "Train[0/50][7/240] Loss: 8.061013221740723, BER : 0.87718359375\n",
            "Train[0/50][8/240] Loss: 7.977132797241211, BER : 0.8683671875\n",
            "Train[0/50][9/240] Loss: 7.974903583526611, BER : 0.868078125\n",
            "Train[0/50][10/240] Loss: 7.9753265380859375, BER : 0.8694609375\n",
            "Train[0/50][11/240] Loss: 8.03271484375, BER : 0.8691484375\n",
            "Train[0/50][12/240] Loss: 8.015864372253418, BER : 0.86117578125\n",
            "Train[0/50][13/240] Loss: 8.188153266906738, BER : 0.85907421875\n",
            "Train[0/50][14/240] Loss: 8.378402709960938, BER : 0.84477734375\n",
            "Train[0/50][15/240] Loss: 8.393698692321777, BER : 0.839109375\n",
            "Train[0/50][16/240] Loss: 8.528434753417969, BER : 0.8401953125\n",
            "Train[0/50][17/240] Loss: 8.586627006530762, BER : 0.838625\n",
            "Train[0/50][18/240] Loss: 8.585226058959961, BER : 0.84272265625\n",
            "Train[0/50][19/240] Loss: 8.6072998046875, BER : 0.839125\n",
            "Train[0/50][20/240] Loss: 8.488789558410645, BER : 0.84205859375\n",
            "Train[0/50][21/240] Loss: 8.415709495544434, BER : 0.8392890625\n",
            "Train[0/50][22/240] Loss: 8.352551460266113, BER : 0.8446953125\n",
            "Train[0/50][23/240] Loss: 8.230137825012207, BER : 0.8392421875\n",
            "Train[0/50][24/240] Loss: 8.243719100952148, BER : 0.84791015625\n",
            "Train[0/50][25/240] Loss: 8.193615913391113, BER : 0.84630078125\n",
            "Train[0/50][26/240] Loss: 8.2012357711792, BER : 0.8520546875\n",
            "Train[0/50][27/240] Loss: 8.148492813110352, BER : 0.84382421875\n",
            "Train[0/50][28/240] Loss: 8.179178237915039, BER : 0.850515625\n",
            "Train[0/50][29/240] Loss: 8.198023796081543, BER : 0.8344921875\n",
            "Train[0/50][30/240] Loss: 8.179018020629883, BER : 0.840125\n",
            "Train[0/50][31/240] Loss: 8.312437057495117, BER : 0.8360390625\n",
            "Train[0/50][32/240] Loss: 8.167531967163086, BER : 0.83917578125\n",
            "Train[0/50][33/240] Loss: 8.112651824951172, BER : 0.83530078125\n",
            "Train[0/50][34/240] Loss: 7.98326301574707, BER : 0.83814453125\n",
            "Train[0/50][35/240] Loss: 7.954312324523926, BER : 0.840515625\n",
            "Train[0/50][36/240] Loss: 7.9321746826171875, BER : 0.83998828125\n",
            "Train[0/50][37/240] Loss: 7.938837051391602, BER : 0.84271484375\n",
            "Train[0/50][38/240] Loss: 7.9776411056518555, BER : 0.835859375\n",
            "Train[0/50][39/240] Loss: 7.91365385055542, BER : 0.83601171875\n",
            "Train[0/50][40/240] Loss: 7.925403594970703, BER : 0.82959375\n",
            "Train[0/50][41/240] Loss: 7.896270751953125, BER : 0.83281640625\n",
            "Train[0/50][42/240] Loss: 7.927243232727051, BER : 0.83244921875\n",
            "Train[0/50][43/240] Loss: 7.871910095214844, BER : 0.82892578125\n",
            "Train[0/50][44/240] Loss: 7.909868240356445, BER : 0.82752734375\n",
            "Train[0/50][45/240] Loss: 7.839835166931152, BER : 0.83039453125\n",
            "Train[0/50][46/240] Loss: 7.8044281005859375, BER : 0.82822265625\n",
            "Train[0/50][47/240] Loss: 7.70984411239624, BER : 0.8275390625\n",
            "Train[0/50][48/240] Loss: 7.719776630401611, BER : 0.82478125\n",
            "Train[0/50][49/240] Loss: 7.647815227508545, BER : 0.824421875\n",
            "Train[0/50][50/240] Loss: 7.592187404632568, BER : 0.8311953125\n",
            "Train[0/50][51/240] Loss: 7.565716743469238, BER : 0.8277890625\n",
            "Train[0/50][52/240] Loss: 7.705008506774902, BER : 0.820859375\n",
            "Train[0/50][53/240] Loss: 7.550854682922363, BER : 0.82262109375\n",
            "Train[0/50][54/240] Loss: 7.682579040527344, BER : 0.8240703125\n",
            "Train[0/50][55/240] Loss: 7.7527971267700195, BER : 0.821640625\n",
            "Train[0/50][56/240] Loss: 7.732005596160889, BER : 0.819546875\n",
            "Train[0/50][57/240] Loss: 7.806859970092773, BER : 0.8148203125\n",
            "Train[0/50][58/240] Loss: 7.791357040405273, BER : 0.81480078125\n",
            "Train[0/50][59/240] Loss: 7.705042839050293, BER : 0.81141015625\n",
            "Train[0/50][60/240] Loss: 7.704450607299805, BER : 0.8136328125\n",
            "Train[0/50][61/240] Loss: 7.6551103591918945, BER : 0.8142578125\n",
            "Train[0/50][62/240] Loss: 7.5417985916137695, BER : 0.8055234375\n",
            "Train[0/50][63/240] Loss: 7.540898323059082, BER : 0.81037890625\n",
            "Train[0/50][64/240] Loss: 7.522022247314453, BER : 0.8126484375\n",
            "Train[0/50][65/240] Loss: 7.47064208984375, BER : 0.808703125\n",
            "Train[0/50][66/240] Loss: 7.420713901519775, BER : 0.80492578125\n"
          ]
        }
      ],
      "source": [
        "checkpoint_path = \"Saved_Trained_Checkpoints/\"\n",
        "Output_Spikes = \"Output_Spikes/\"\n",
        "Enc_syn_Spikes = \"Enc_syn_Spikes/\"\n",
        "Intermediate_Lyrs = \"Intermediate_Lyrs/\"\n",
        "epoch_activations_list = [] # Create a list to store activations for each epoch\n",
        "epoch_activations = {}\n",
        "\n",
        "\n",
        "# ///////////////////////////////////\n",
        "\n",
        "# Define hook_layers and hook_names\n",
        "hook_layers = [net.encoder[2], net.encoder[5], net.encoder[8], net.encoder[11], net.decoder[1], net.decoder[4], net.decoder[7], net.decoder[9]]\n",
        "hook_names = [\"Enc_Lk1\", \"Enc_syn1\", \"Enc_syn2\", \"Enc_Lk2\", \"Dec_Lk1\", \"Dec_syn1\", \"Dec_syn2\", \"Dec_Lk2\"]\n",
        "\n",
        "# Create an empty dictionary to store activations\n",
        "epoch_activations = {}\n",
        "\n",
        "# Register hooks for capturing activations\n",
        "hooks = []\n",
        "for i, layer in enumerate(hook_layers):\n",
        "    hook_fn = get_activation(hook_names[i])\n",
        "    hooks.append(layer.register_forward_hook(hook_fn))\n",
        "\n",
        "# Run training and testing\n",
        "for e in range(epochs):\n",
        "    train_loss = train(net, train_loader, optimizer, e)\n",
        "    train_avg_loss_rec.append(sum(train_loss_rec) / len(train_loader))\n",
        "\n",
        "    test_loss = test(net, test_loader, optimizer, e)\n",
        "    test_avg_loss_rec.append(sum(test_loss_rec) / (len(test_loader)))\n",
        "\n",
        "    # #---------------------------------------------------------- Access the out_en tensor\n",
        "    out_en = net.out_en\n",
        "    out_en_numpy = out_en.cpu().detach().numpy()\n",
        "    # Save with a different name for each epoch\n",
        "    out_en_filename = Output_Spikes + f\"out_en_epoch_{e + 1}.npy\"\n",
        "    np.save(out_en_filename, out_en_numpy)\n",
        "\n",
        "    # #-----------------------------------------------------------Access the out tensor\n",
        "    out = net.out\n",
        "    out_numpy = out.cpu().detach().numpy()\n",
        "    # Save with a different name for each epoch\n",
        "    out_filename = Output_Spikes + f\"out_epoch_{e + 1}.npy\"\n",
        "    np.save(out_filename, out_numpy)\n",
        "    # -------------------------------------------------------------------- Intermediate Layers\n",
        "# Check if the current epoch is a multiple of 10\n",
        "    if (e + 1) % 10 == 0:\n",
        "        # Save the epoch_activations dictionary to a file\n",
        "        activations_path = Intermediate_Lyrs +  f\"activations_epoch_{e + 1}.pkl\"\n",
        "        with open(activations_path, 'wb') as file:\n",
        "            pickle.dump(epoch_activations, file)\n",
        "    # Capture activations for the current epoch\n",
        "    epoch_activations[e] = {}\n",
        "    for i, name in enumerate(hook_names):\n",
        "        epoch_activations[e][name] = activation.get(name, None)  # Use get to avoid KeyError\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "# beta_syn=0.0001\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "1sKLRccI0er_"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}