notebooks/analog_training_LeNet5_hwa.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LeNet5 Analog Hardware-aware Training Example\n",
    " Training the LeNet5 neural network with hardware aware training using the Inference RPU Config to optimize inference on a PCM device.\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/IBM/aihwkit/blob/master/notebooks//analog_training_LeNet5_hwa.ipynb\" target=\"_parent\">\n",
    "    <img src=\"https://colab.research.google.com/assets/colab-badge.svg\"/>\n",
    "</a>\n",
    "\n",
    "## IBM Analog Hardware Acceleration Kit\n",
    "\n",
    "IBM Analog Hardware Acceleration Kit (AIHWKIT) is an open source Python toolkit for exploring and using the capabilities of in-memory computing devices in the context of artificial intelligence.\n",
    "The pytorch integration consists of a series of primitives and features that allow using the toolkit within PyTorch. \n",
    "The github repository can be found at: https://github.com/IBM/aihwkit\n",
    "\n",
    "There are two possible scenarios for using Analog AI, one where the Analog accelerator targets training of DNN and one where the Analog accelerator aims at accelerating the inference of a DNN.\n",
    "Employing Analog accelerator for training scenarios requires innovations on the algorithms used during the backpropagation (BP) phase and we will see an example of such innovation in the Tiki-Taka notebook.\n",
    "Employing Analog accelerator for inference scenarios allows the use of a digital accelerator for the training part and then transfering the weights to the analog hardware for the inference, which we will explore in this notebook.\n",
    "\n",
    "\n",
    "## Hardware Aware Training with Analog AI\n",
    "\n",
    "When performing inference on Analog hardware, direct transfer of the weights trained on a digital chip to analog hardware would results in reduced network accuracies due to non-idealities of the non-volatile memory technologies used in the analog chip. \n",
    "For example, a common non-ideality of Phase Change Memory is the resistance drift. This drift is due to structural relaxation of the physical material composing the memory and causes an evolution of the memory resistance over time, which translates in a change in the weight stored in the memory array. \n",
    "This would eventually result in a decrease network accuracy over time.\n",
    "What one can do is to train the network on digital accelerator but in a way that is aware of the hardware characteristics that will be used in the inference pass, we refer to this as Hardware Aware Training (HWA).\n",
    "\n",
    "<center><img src=\"img/hwa.jpg\" style=\"width:50%; height:50%\"/></center> \n",
    "\n",
    "In hardware-aware training, we add many of the device non-idealities in the forward pass, so that the network itself, while learning the features to achieve high accuracy, also learns how to be resiliant to these non-idealities.\n",
    "Examples of the non-idealities inserted during the forward pass include quantization noise and thermal noise related to the Digital-to-Analog converter (DAC) and Analog-to-Digital converter (ADC) for the digital-to-analog signal conversion (and vice versa), as well as non-idealities of the NVM technology in use, which for Phase Change Memory are the state dependent weight noise, the weight drift among others.\n",
    "\n",
    "Hardware-aware training greatly improves analog inference accuracy. For more information, please refer to:\n",
    "\n",
    "https://ieeexplore.ieee.org/document/8993573\n",
    "\n",
    "https://ieeexplore.ieee.org/document/8776519\n",
    "\n",
    "https://www.nature.com/articles/s41467-020-16108-9\n",
    "\n",
    "In this notebook we will use the AIHWKIT to do HWA training of a LeNet5 inspired network with the MNIST dataset. We will then evaluate the accuracy of the analog hardware and its evolution in time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first thing to do is to install the AIHKIT and dependencies in your environment. The preferred way to install this package is by using the Python package index (please uncomment this line to install in your environment if not previously installed):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To install the cpu-only enabled kit, un-comment the line below\n",
    "# %pip install aihwkit\n",
    "\n",
    "# To install the GPU-enabled wheel go to https://aihwkit.readthedocs.io/en/latest/advanced_install.html#install-the-aihwkit-using-pip\n",
    "# and copy the option on GPU options that best suits your enviroment and paste it below and run the cell. For example, Python 3.10 and CUDA 12.1:\n",
    "# !wget https://aihwkit-gpu-demo.s3.us-east.cloud-object-storage.appdomain.cloud/aihwkit-0.9.2+cuda121-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\n",
    "# %pip install aihwkit-0.9.2+cuda121-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the library was installed correctly, you can use the following snippet for creating an analog layer and predicting the output:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import Tensor\n",
    "from aihwkit.nn import AnalogLinear\n",
    "\n",
    "model = AnalogLinear(2, 2)\n",
    "model(Tensor([[0.1, 0.2], [0.3, 0.4]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the package is installed and running, we can start working on creating the LeNet5 network.\n",
    "\n",
    "AIHWKIT offers different Analog layers that can be used to build a network, including AnalogLinear and AnalogConv2d which will be the main layers used to build the present network. \n",
    "In addition to the standard input that are expected by the PyTorch layers (in_channels, out_channels, etc.) the analog layers also expect a rpu_config input which defines various settings of the RPU tile. Through the rpu_config parameter the user can specify many of the hardware specs such as: device used in the cross-point array, bit used by the ADC/DAC converters, noise values and many other. Additional details on the RPU configuration can be found at https://aihwkit.readthedocs.io/en/latest/using_simulator.html#rpu-configurations\n",
    "For this particular case we will use two device per cross-point which will effectively allow us to enable the weight transfer needed to implement the Tiki-Taka algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_rpu_config():\n",
    "\n",
    "    from aihwkit.simulator.configs import InferenceRPUConfig\n",
    "    from aihwkit.simulator.configs.utils import BoundManagementType, WeightNoiseType\n",
    "    from aihwkit.inference import PCMLikeNoiseModel\n",
    "\n",
    "    rpu_config = InferenceRPUConfig()\n",
    "    rpu_config.backward.bound_management = BoundManagementType.NONE\n",
    "    rpu_config.forward.out_res = -1.  # Turn off (output) ADC discretization.\n",
    "    rpu_config.forward.w_noise_type = WeightNoiseType.ADDITIVE_CONSTANT\n",
    "    rpu_config.forward.w_noise = 0.02\n",
    "    rpu_config.noise_model = PCMLikeNoiseModel(g_max=25.0)\n",
    "\n",
    "    return rpu_config"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now use this rpu_config as input of the network model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn import Tanh, MaxPool2d, LogSoftmax, Flatten\n",
    "from aihwkit.nn import AnalogConv2d, AnalogLinear, AnalogSequential\n",
    "\n",
    "def create_analog_network(rpu_config):\n",
    "    \n",
    "    channel = [16, 32, 512, 128]\n",
    "    model = AnalogSequential(\n",
    "        AnalogConv2d(in_channels=1, out_channels=channel[0], kernel_size=5, stride=1,\n",
    "                        rpu_config=rpu_config),\n",
    "        Tanh(),\n",
    "        MaxPool2d(kernel_size=2),\n",
    "        AnalogConv2d(in_channels=channel[0], out_channels=channel[1], kernel_size=5, stride=1,\n",
    "                        rpu_config=rpu_config),\n",
    "        Tanh(),\n",
    "        MaxPool2d(kernel_size=2),\n",
    "        Tanh(),\n",
    "        Flatten(),\n",
    "        AnalogLinear(in_features=channel[2], out_features=channel[3], rpu_config=rpu_config),\n",
    "        Tanh(),\n",
    "        AnalogLinear(in_features=channel[3], out_features=10, rpu_config=rpu_config),\n",
    "        LogSoftmax(dim=1)\n",
    "    )\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use the cross entropy to calculate the loss and the Stochastic Gradient Descent (SGD) as optimizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn import CrossEntropyLoss\n",
    "\n",
    "criterion = CrossEntropyLoss()\n",
    "\n",
    "\n",
    "from aihwkit.optim import AnalogSGD\n",
    "\n",
    "def create_analog_optimizer(model):\n",
    "    \"\"\"Create the analog-aware optimizer.\n",
    "\n",
    "    Args:\n",
    "        model (nn.Module): model to be trained\n",
    "\n",
    "    Returns:\n",
    "        Optimizer: created analog optimizer\n",
    "    \"\"\"\n",
    "    \n",
    "    optimizer = AnalogSGD(model.parameters(), lr=0.01) # we will use a learning rate of 0.01 as in the paper\n",
    "    optimizer.regroup_param_groups(model)\n",
    "\n",
    "    return optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now write the train function which will optimize the network over the MNIST train dataset. The train_step function will take as input the images to train on, the model to train and the criterion and optimizer to train with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import device, cuda\n",
    "\n",
    "DEVICE = device('cuda' if cuda.is_available() else 'cpu')\n",
    "print('Running the simulation on: ', DEVICE)\n",
    "\n",
    "def train_step(train_data, model, criterion, optimizer):\n",
    "    \"\"\"Train network.\n",
    "\n",
    "    Args:\n",
    "        train_data (DataLoader): Validation set to perform the evaluation\n",
    "        model (nn.Module): Trained model to be evaluated\n",
    "        criterion (nn.CrossEntropyLoss): criterion to compute loss\n",
    "        optimizer (Optimizer): analog model optimizer\n",
    "\n",
    "    Returns:\n",
    "        train_dataset_loss: epoch loss of the train dataset\n",
    "    \"\"\"\n",
    "    total_loss = 0\n",
    "\n",
    "    model.train()\n",
    "\n",
    "    for images, labels in train_data:\n",
    "        images = images.to(DEVICE)\n",
    "        labels = labels.to(DEVICE)\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # Add training Tensor to the model (input).\n",
    "        output = model(images)\n",
    "        loss = criterion(output, labels)\n",
    "\n",
    "        # Run training (backward propagation).\n",
    "        loss.backward()\n",
    "\n",
    "        # Optimize weights.\n",
    "        optimizer.step()\n",
    "        total_loss += loss.item() * images.size(0)\n",
    "    train_dataset_loss = total_loss / len(train_data.dataset)\n",
    "\n",
    "    return train_dataset_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since training can be quite time consuming it is nice to see the evolution of the training process by testing the model capabilities on a set of images that it has not seen before (test dataset). So we write a test_step function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_step(validation_data, model, criterion):\n",
    "    \"\"\"Test trained network\n",
    "\n",
    "    Args:\n",
    "        validation_data (DataLoader): Validation set to perform the evaluation\n",
    "        model (nn.Module): Trained model to be evaluated\n",
    "        criterion (nn.CrossEntropyLoss): criterion to compute loss\n",
    "\n",
    "    Returns: \n",
    "        test_dataset_loss: epoch loss of the train_dataset\n",
    "        test_dataset_error: error of the test dataset\n",
    "        test_dataset_accuracy: accuracy of the test dataset\n",
    "    \"\"\"\n",
    "    total_loss = 0\n",
    "    predicted_ok = 0\n",
    "    total_images = 0\n",
    "\n",
    "    model.eval()\n",
    "\n",
    "    for images, labels in validation_data:\n",
    "        images = images.to(DEVICE)\n",
    "        labels = labels.to(DEVICE)\n",
    "\n",
    "        pred = model(images)\n",
    "        loss = criterion(pred, labels)\n",
    "        total_loss += loss.item() * images.size(0)\n",
    "\n",
    "        _, predicted = torch.max(pred.data, 1)\n",
    "        total_images += labels.size(0)\n",
    "        predicted_ok += (predicted == labels).sum().item()\n",
    "        test_dataset_accuracy = predicted_ok/total_images*100\n",
    "        test_dataset_error = (1-predicted_ok/total_images)*100\n",
    "\n",
    "    test_dataset_loss = total_loss / len(validation_data.dataset)\n",
    "\n",
    "    return test_dataset_loss, test_dataset_error, test_dataset_accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To reach satisfactory accuracy levels, the train_step will have to be repeated mulitple time so we will implement a loop over a certain number of epochs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def training_loop(model, criterion, optimizer, train_data, validation_data, epochs=15, print_every=1):\n",
    "    \"\"\"Training loop.\n",
    "\n",
    "    Args:\n",
    "        model (nn.Module): Trained model to be evaluated\n",
    "        criterion (nn.CrossEntropyLoss): criterion to compute loss\n",
    "        optimizer (Optimizer): analog model optimizer\n",
    "        train_data (DataLoader): Validation set to perform the evaluation\n",
    "        validation_data (DataLoader): Validation set to perform the evaluation\n",
    "        epochs (int): global parameter to define epochs number\n",
    "        print_every (int): defines how many times to print training progress\n",
    "\n",
    "    \"\"\"\n",
    "    train_losses = []\n",
    "    valid_losses = []\n",
    "    test_error = []\n",
    "\n",
    "    # Train model\n",
    "    for epoch in range(0, epochs):\n",
    "        # Train_step\n",
    "        train_loss = train_step(train_data, model, criterion, optimizer)\n",
    "        train_losses.append(train_loss)\n",
    "\n",
    "        if epoch % print_every == (print_every - 1):\n",
    "            # Validate_step\n",
    "            with torch.no_grad():\n",
    "                valid_loss, error, accuracy = test_step(validation_data, model, criterion)\n",
    "                valid_losses.append(valid_loss)\n",
    "                test_error.append(error)\n",
    "\n",
    "            print(f'Epoch: {epoch}\\t'\n",
    "                  f'Train loss: {train_loss:.4f}\\t'\n",
    "                  f'Valid loss: {valid_loss:.4f}\\t'\n",
    "                  f'Test error: {error:.2f}%\\t'\n",
    "                  f'Test accuracy: {accuracy:.2f}%\\t')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_inference(model, criterion, test_data):\n",
    "    \n",
    "    from numpy import logspace, log10\n",
    "    \n",
    "    total_loss = 0\n",
    "    predicted_ok = 0\n",
    "    total_images = 0\n",
    "    accuracy_pre = 0\n",
    "    error_pre = 0\n",
    "    \n",
    "    # Create the t_inference_list using inference_time.\n",
    "    # Generate the 9 values between 0 and the inference time using log10\n",
    "    max_inference_time = 1e6\n",
    "    n_times = 9\n",
    "    t_inference_list = [0.0] + logspace(0, log10(float(max_inference_time)), n_times).tolist()\n",
    "\n",
    "    # Simulation of inference pass at different times after training.\n",
    "    for t_inference in t_inference_list:\n",
    "        model.drift_analog_weights(t_inference)\n",
    "\n",
    "        time_since = t_inference\n",
    "        accuracy_post = 0\n",
    "        error_post = 0\n",
    "        predicted_ok = 0\n",
    "        total_images = 0\n",
    "\n",
    "        for images, labels in test_data:\n",
    "            images = images.to(DEVICE)\n",
    "            labels = labels.to(DEVICE)\n",
    "\n",
    "            pred = model(images)\n",
    "            loss = criterion(pred, labels)\n",
    "            total_loss += loss.item() * images.size(0)\n",
    "\n",
    "            _, predicted = torch.max(pred.data, 1)\n",
    "            total_images += labels.size(0)\n",
    "            predicted_ok += (predicted == labels).sum().item()\n",
    "            accuracy_post = predicted_ok/total_images*100\n",
    "            error_post = (1-predicted_ok/total_images)*100\n",
    "\n",
    "        print(f'Error after inference: {error_post:.2f}\\t'\n",
    "              f'Accuracy after inference: {accuracy_post:.2f}%\\t'\n",
    "              f'Drift t={time_since: .2e}\\t')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now download the MNIST dataset and prepare the images for the training and test:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from torchvision import datasets, transforms\n",
    "PATH_DATASET = os.path.join('data', 'DATASET')\n",
    "os.makedirs(PATH_DATASET, exist_ok=True)\n",
    "\n",
    "def load_images():\n",
    "    \"\"\"Load images for train from torchvision datasets.\"\"\"\n",
    "\n",
    "    transform = transforms.Compose([transforms.ToTensor()])\n",
    "    train_set = datasets.MNIST(PATH_DATASET, download=True, train=True, transform=transform)\n",
    "    test_set = datasets.MNIST(PATH_DATASET, download=True, train=False, transform=transform)\n",
    "    train_data = torch.utils.data.DataLoader(train_set, batch_size=8, shuffle=True)\n",
    "    test_data = torch.utils.data.DataLoader(test_set, batch_size=8, shuffle=False)\n",
    "\n",
    "    return train_data, test_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Put together all the code above to train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "torch.manual_seed(1)\n",
    "\n",
    "#load the dataset\n",
    "train_data, test_data = load_images()\n",
    "\n",
    "#create the rpu_config\n",
    "rpu_config = create_rpu_config()\n",
    "\n",
    "#create the model\n",
    "model = create_analog_network(rpu_config).to(DEVICE)\n",
    "\n",
    "#define the analog optimizer\n",
    "optimizer = create_analog_optimizer(model)\n",
    "\n",
    "training_loop(model, criterion, optimizer, train_data, test_data)\n",
    "\n",
    "print('\\nStarting Inference Pass:\\n')\n",
    "\n",
    "test_inference(model, criterion, test_data)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "8afadb82c8c635d284d204a78cd7f3b56094702ee8f92f25084bfbbc5b27362b"
  },
  "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}