docs/_static/AnalogVNN_Demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AnalogVNN Demo/Tutorial\n",
    "\n",
    "Copyright © 2021-present Vivswan Shah (vivswanshah@pitt.edu)\n",
    "\n",
    "<br>\n",
    "\n",
    "[![arXiv](https://img.shields.io/badge/arXiv-2210.10048-orange.svg)](https://arxiv.org/abs/2210.10048)\n",
    "[![PyPI version](https://badge.fury.io/py/analogvnn.svg)](https://badge.fury.io/py/analogvnn)\n",
    "[![Documentation Status](https://readthedocs.org/projects/analogvnn/badge/?version=stable)](https://analogvnn.readthedocs.io/en/stable/?badge=stable)\n",
    "[![Python](https://img.shields.io/badge/python-3.7--3.10-blue)](https://badge.fury.io/py/analogvnn)\n",
    "[![License: MPL 2.0](https://img.shields.io/badge/License-MPL_2.0-blue.svg)](https://opensource.org/licenses/MPL-2.0)\n",
    "\n",
    "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://analogvnn.readthedocs.io/en/release/tutorial.html\">\n",
    "      <center>\n",
    "        <img src=\"https://analogvnn.readthedocs.io/en/in_progess/_static/analogvnn-logo-square-black.svg\" height=\"32px\" />\n",
    "      </center>\n",
    "      View on AnalogVNN\n",
    "    </a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/Vivswan/AnalogVNN/blob/release/docs/_static/AnalogVNN_Demo.ipynb\">\n",
    "      <center>\n",
    "        <img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />\n",
    "      </center>\n",
    "      Run in Google Colab\n",
    "    </a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://github.com/Vivswan/AnalogVNN/blob/release/docs/_static/AnalogVNN_Demo.ipynb\">\n",
    "      <center>\n",
    "        <img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />\n",
    "      </center>\n",
    "      View source on GitHub\n",
    "    </a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://github.com/Vivswan/AnalogVNN/raw/release/docs/_static/AnalogVNN_Demo.ipynb\">\n",
    "      <center>\n",
    "        <img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />\n",
    "      </center>\n",
    "      Download notebook\n",
    "    </a>\n",
    "  </td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### To create 3 layered linear photonic analog neural network with 4-bit [precision](https://analogvnn.readthedocs.io/en/release/extra_classes.html#reduceprecision), 0.5 [leakage](https://analogvnn.readthedocs.io/en/release/extra_classes.html#leakage-or-error-probability) and [clamp](https://analogvnn.readthedocs.io/en/release/extra_classes.html#clamp) normalization:\n",
    "\n",
    "![3 Layered Linear Photonic Analog Neural Network](analogvnn_model.png)\n",
    "\n",
    "Python file:\n",
    "[Sample code](https://github.com/Vivswan/AnalogVNN/blob/release/sample_code.py)\n",
    "and\n",
    "[Sample code with logs](https://github.com/Vivswan/AnalogVNN/blob/release/sample_code_with_logs.py)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Setting up the Enviroment AnalogVNN"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "812kuN10TZgu"
   },
   "outputs": [],
   "source": [
    "# Install AnalogVNN with Pip\n",
    "!pip install analogvnn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0trJmd6DjqBZ"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import analogvnn\n",
    "\n",
    "import torch.backends.cudnn\n",
    "import torchvision\n",
    "from torch import optim, nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision.transforms import transforms\n",
    "\n",
    "from analogvnn.nn.Linear import Linear\n",
    "from analogvnn.nn.activation.Gaussian import GeLU\n",
    "from analogvnn.nn.module.FullSequential import FullSequential\n",
    "from analogvnn.nn.noise.GaussianNoise import GaussianNoise\n",
    "from analogvnn.nn.normalize.Clamp import Clamp\n",
    "from analogvnn.nn.precision.ReducePrecision import ReducePrecision\n",
    "from analogvnn.parameter.PseudoParameter import PseudoParameter\n",
    "from analogvnn.utils.is_cpu_cuda import is_cpu_cuda\n",
    "\n",
    "print(f\"AnalogVNN version: {analogvnn.__version__}\")\n",
    "print(f\"PyTorch version: {torch.__version__}\")\n",
    "\n",
    "is_cpu_cuda.use_cuda_if_available()\n",
    "torch.manual_seed(0)\n",
    "device, is_cuda = is_cpu_cuda.is_using_cuda\n",
    "print(f\"Device: {device}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7NAbSZiaoJ4z"
   },
   "source": [
    "## Load a dataset\n",
    "\n",
    "Load and prepare the [MNIST dataset](http://yann.lecun.com/exdb/mnist/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7FP5258xjs-v"
   },
   "outputs": [],
   "source": [
    "def load_vision_dataset(dataset, path, batch_size, is_cuda=False, grayscale=True):\n",
    "    dataset_kwargs = {\n",
    "        'batch_size': batch_size,\n",
    "        'shuffle': True\n",
    "    }\n",
    "\n",
    "    if is_cuda:\n",
    "        cuda_kwargs = {\n",
    "            'num_workers': 1,\n",
    "            'pin_memory': True,\n",
    "        }\n",
    "        dataset_kwargs.update(cuda_kwargs)\n",
    "\n",
    "    if grayscale:\n",
    "        use_transform = transforms.Compose([\n",
    "            transforms.Grayscale(),\n",
    "            transforms.ToTensor(),\n",
    "        ])\n",
    "    else:\n",
    "        use_transform = transforms.Compose([transforms.ToTensor()])\n",
    "\n",
    "    train_set = dataset(root=path, train=True, download=True, transform=use_transform)\n",
    "    test_set = dataset(root=path, train=False, download=True, transform=use_transform)\n",
    "    train_loader = DataLoader(train_set, **dataset_kwargs)\n",
    "    test_loader = DataLoader(test_set, **dataset_kwargs)\n",
    "\n",
    "    zeroth_element = next(iter(test_loader))[0]\n",
    "    input_shape = list(zeroth_element.shape)\n",
    "\n",
    "    return train_loader, test_loader, input_shape, tuple(train_set.classes)\n",
    "\n",
    "# Loading Data...\n",
    "print(f\"Loading Data...\")\n",
    "train_loader, test_loader, input_shape, classes = load_vision_dataset(\n",
    "    dataset=torchvision.datasets.MNIST,\n",
    "    path=\"_data/\",\n",
    "    batch_size=128,\n",
    "    is_cuda=is_cuda\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "06PdU_iGVe16"
   },
   "source": [
    "## Build a 3 layered linear photonic analog neural network\n",
    "\n",
    "[`FullSequential`](https://analogvnn.readthedocs.io/en/release/autoapi/analogvnn/nn/module/FullSequential/index.html#analogvnn.nn.module.FullSequential.FullSequential) is sequential model where backward graph is the reverse of forward graph.\n",
    "\n",
    "To add the [Reduce Precision](https://analogvnn.readthedocs.io/en/release/extra_classes.html#reduce-precision), [Normalization](https://analogvnn.readthedocs.io/en/release/extra_classes.html#normalization), and [Noise](https://analogvnn.readthedocs.io/en/release/extra_classes.html#noise) before and after the main Linear layer, `add_layer` function is used.\n",
    "\n",
    "Leakage definition: [https://analogvnn.readthedocs.io/en/release/extra_classes.html#leakage-or-error-probability](https://analogvnn.readthedocs.io/en/release/extra_classes.html#leakage-or-error-probability)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "h3IKyzTCDNGo"
   },
   "outputs": [],
   "source": [
    "class LinearModel(FullSequential):\n",
    "    def __init__(self, activation_class, norm_class, precision_class, precision, noise_class, leakage):\n",
    "        super().__init__()\n",
    "\n",
    "        self.activation_class = activation_class\n",
    "        self.norm_class = norm_class\n",
    "        self.precision_class = precision_class\n",
    "        self.precision = precision\n",
    "        self.noise_class = noise_class\n",
    "        self.leakage = leakage\n",
    "\n",
    "        self.all_layers = []\n",
    "        self.all_layers.append(nn.Flatten(start_dim=1))\n",
    "        self.add_layer(Linear(in_features=28 * 28, out_features=256))\n",
    "        self.add_layer(Linear(in_features=256, out_features=128))\n",
    "        self.add_layer(Linear(in_features=128, out_features=10))\n",
    "\n",
    "        self.add_sequence(*self.all_layers)\n",
    "\n",
    "    def add_layer(self, layer):\n",
    "        self.all_layers.append(self.norm_class())\n",
    "        self.all_layers.append(self.precision_class(precision=self.precision))\n",
    "        self.all_layers.append(self.noise_class(leakage=self.leakage, precision=self.precision))\n",
    "        self.all_layers.append(layer)\n",
    "        self.all_layers.append(self.noise_class(leakage=self.leakage, precision=self.precision))\n",
    "        self.all_layers.append(self.norm_class())\n",
    "        self.all_layers.append(self.precision_class(precision=self.precision))\n",
    "        self.all_layers.append(self.activation_class())\n",
    "        self.activation_class.initialise_(layer.weight)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iOkIKXWoZbmn"
   },
   "source": [
    "Note: [`analogvnn.nn.module.Sequential.Sequential.add_sequence()`](https://analogvnn.readthedocs.io/en/release/autoapi/analogvnn/nn/module/Sequential/index.html#analogvnn.nn.module.Sequential.Sequential.add_sequence) is used to create and set forward and backward graphs in AnalogVNN, more information in Inner Workings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "39SIO0ICWiJ2"
   },
   "outputs": [],
   "source": [
    "nn_model = LinearModel(\n",
    "    activation_class=GeLU,\n",
    "    norm_class=Clamp,\n",
    "    precision_class=ReducePrecision,\n",
    "    precision=2 ** 4,\n",
    "    noise_class=GaussianNoise,\n",
    "    leakage=0.5\n",
    ")\n",
    "print(nn_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BPZ68wASog_I"
   },
   "source": [
    "## Build a WeightModel \n",
    "\n",
    "WeightModel is used to parametrize the parameter of LinearModel to simulate photonic weights\n",
    "\n",
    "[`FullSequential`](https://analogvnn.readthedocs.io/en/release/autoapi/analogvnn/nn/module/FullSequential/index.html#analogvnn.nn.module.FullSequential.FullSequential) is sequential model where backward graph is the reverse of forward graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tcN3WT8zWXQv"
   },
   "outputs": [],
   "source": [
    "class WeightModel(FullSequential):\n",
    "    def __init__(self, norm_class, precision_class, precision, noise_class, leakage):\n",
    "        super().__init__()\n",
    "        self.all_layers = []\n",
    "\n",
    "        self.all_layers.append(norm_class())\n",
    "        self.all_layers.append(precision_class(precision=precision))\n",
    "        self.all_layers.append(noise_class(leakage=leakage, precision=precision))\n",
    "\n",
    "        self.eval()\n",
    "        self.add_sequence(*self.all_layers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Dsudt6dXZBnV"
   },
   "source": [
    "Note: Since the `WeightModel` will only be used for converting the data to analog data to be used in the main `LinearModel`, we can use `eval()` to make sure the `WeightModel` is never been trained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FgehAu7qWlyV"
   },
   "outputs": [],
   "source": [
    "weight_model = WeightModel(\n",
    "    norm_class=Clamp,\n",
    "    precision_class=ReducePrecision,\n",
    "    precision=2 ** 4,\n",
    "    noise_class=GaussianNoise,\n",
    "    leakage=0.5\n",
    ")\n",
    "print(weight_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Dtg27Y80WwR0"
   },
   "source": [
    "Using [`PseudoParameter`](https://analogvnn.readthedocs.io/en/release/inner_workings.html#pseudoparameters) to parametrize the parameter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "O8i_yEZHWpZb"
   },
   "outputs": [],
   "source": [
    "PseudoParameter.parametrize_module(nn_model, transformation=weight_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ix4mEL65on-w"
   },
   "source": [
    "## Adding accuracy, loss and optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bL9owooIXaJ2"
   },
   "outputs": [],
   "source": [
    "def cross_entropy_accuracy(output, target) -> float:\n",
    "    _, preds = torch.max(output.data, 1)\n",
    "    correct = (preds == target).sum().item()\n",
    "    return correct / len(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XnWd-loMXfLj"
   },
   "outputs": [],
   "source": [
    "nn_model.loss_function = nn.CrossEntropyLoss()\n",
    "nn_model.accuracy_function = cross_entropy_accuracy\n",
    "nn_model.optimizer = optim.Adam(params=nn_model.parameters())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "txh1-HanXkzC"
   },
   "source": [
    "## Compiling the Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wbmpzLWpXo5t"
   },
   "outputs": [],
   "source": [
    "nn_model.compile(device=device)\n",
    "weight_model.compile(device=device)\n",
    "print(\"Compiled\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F0BBx6D4XExj"
   },
   "source": [
    "## Train and evaluate your model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9cAdBEaXX51G"
   },
   "outputs": [],
   "source": [
    "for epoch in range(10):\n",
    "    train_loss, train_accuracy = nn_model.train_on(train_loader, epoch=epoch)\n",
    "    test_loss, test_accuracy = nn_model.test_on(test_loader, epoch=epoch)\n",
    "\n",
    "    str_epoch = str(epoch + 1).zfill(1)\n",
    "    print_str = f'({str_epoch})' \\\n",
    "                f' Training Loss: {train_loss:.4f},' \\\n",
    "                f' Training Accuracy: {100. * train_accuracy:.0f}%,' \\\n",
    "                f' Testing Loss: {test_loss:.4f},' \\\n",
    "                f' Testing Accuracy: {100. * test_accuracy:.0f}%\\n'\n",
    "    print(print_str)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-47O6_GLdRuT"
   },
   "source": [
    "## Conclusion\n",
    "\n",
    "Congratulations! You have trained a 3 layered linear photonic analog neural network with 4-bit [precision](https://analogvnn.readthedocs.io/en/release/extra_classes.html#reduceprecision), 0.5 [leakage](https://analogvnn.readthedocs.io/en/release/extra_classes.html#leakage-or-error-probability) and [clamp](https://analogvnn.readthedocs.io/en/release/extra_classes.html#clamp) normalization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gic6z7KcYo1h"
   },
   "source": [
    "GitHub: [https://github.com/Vivswan/AnalogVNN](https://github.com/Vivswan/AnalogVNN)\n",
    "\n",
    "Documentation: [https://analogvnn.readthedocs.io/](https://analogvnn.readthedocs.io/)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}