notebooks/analog_training_LeNet5.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vjqem3CNzHeP"
   },
   "source": [
    "# LeNet5 Analog Training Example\n",
    "Training the LeNet5 neural network with MNIST dataset and the Analog SGD optimizer simulated on the analog resistive random-access memory with soft bounds (ReRam) device.\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/IBM/aihwkit/blob/master/notebooks/analog_training_LeNet5.ipynb\" target=\"_parent\">\n",
    "    <img src=\"https://colab.research.google.com/assets/colab-badge.svg\"/>\n",
    "</a>\n",
    "\n",
    "## Why Analog AI\n",
    "\n",
    "In-memory computing hardware increases the speed and energy-efficiency needed for the next steps in AI. Analog AI delivers radical performance improvements by combining compute and memory in a single device, eliminating the von Neumann bottleneck. \n",
    "Based on von Neumann architecture, conventional computers perform calculations by repeatedly transferring data between the memory and processor. These trips require time and energy, negatively impacting performance. This is known as the von Neumann bottleneck. \n",
    "\n",
    "<img src=\"img/processing-unit-and-computional-memory.png\"  alt=\"Drawing\" style=\"width=50px;\"/> \n",
    "<img src=\"img/processing-unit-and-conventional-memory.png\" alt=\"Drawing\" style=\"width=50px;\"> \n",
    "\n",
    "By leveraging the physical properties of in-memory computing devices (example: Phase Change Memory or PCM), computation happens at the same place where the data is stored, drastically reducing energy consumption. Because there is no movement of data, tasks can be performed in a fraction of the time and with much less energy. This is different from a conventional computer, where the data is transferred from the DRAM memory to the CPU every time a computation is done. For example, moving 64 bits of data from DRAM to CPU consumes 1-2nJ, which is 10,000–2,000,000 times more energy than is dissipated in a PCM device performing a multiplication operation (1-100fJ). Also, PCM does not consume power when the devices are inactive, and the data will be retained for up to 10 years even when the power supply is turned off. \n",
    "\n",
    "## The physics behind PCM\n",
    "\n",
    " With PCM, when an electrical pulse is applied to the material, it changes the [conductance](https://energyeducation.ca/encyclopedia/Electrical_conductance) of the device by switching the material between amorphous and crystalline phases. A low electrical pulse will make the PCM device more crystalline (less resistance), this pulse can be repeatedly applied to gradually decrease the device resistance. On the other hand, the material change from the crystalline phase (low resistance) to the amorphouse phase (high resistance) is quite abrupt and requires a high electrical pulse to RESET the device. Therefore, it is possible to records the states as a continuum of values between the two extremes, instead of encoding 0 or 1 like in the digital world.\n",
    "\n",
    "PCM devices have the ability to store synaptic weights in their analog conductance state. When PCM devices are arranged in a crossbar configuration, it allows to perform an analog matrix-vector multiplication in a single time step, exploiting the advantages of multi-level storage capability and Kirchhoff’s circuits laws. The figure below shows how PCM devices are arranged in a crossbar configuration. \n",
    "<center><img src=\"img/pcm-array.png\" style=\"width:30%; height:30%\"/></center> \n",
    "\n",
    "This crossbar configuration is also referred to as an Analog tile. The PCM devices at each crossbard crosspoint are also referred to as Resitive Processing Units or RPU units as shown in the figure below:\n",
    "\n",
    "<center><img src=\"img/pcm_rpu_unit.png\" style=\"width:30%; height:30%\"/></center> \n",
    "\n",
    "Besides PCM, other devices or materials can be used as resistive units or RPUs in the crossbar configuration. Examples include Resistive Random Access Memory (RRAM), Electro chemical Random Access Memory (ECRAM), Magnetic RAM (MRAM), photonics, etc. The weights of neural network are stored in the RPU units as confuctance values that are programmed on the chip through as series of electrical pulses. The conductance behavior changes from one analog device to another. \n",
    "\n",
    "\n",
    "## Analog AI and Neural Networks\n",
    "\n",
    "In deep learning inference, data propagation through multiple layers of a neural network involves a sequence of matrix multiplications, as each layer can be represented as a matrix of synaptic weights. On an Analog chip, these weights are stored in the conductance states of resistive devices such as PCM. The devices are arranged in crossbar arrays, creating an artificial neural network where all matrix multiplications are performed in-place in an analog manner. This structure allows inference to be performed using little energy with high areal density of synapses. An in-memory computing chip typically consists of multiple crossbar arrays of memory devices that communicate with each other (see figure below). A neural network layer can be implemented on (at least) one crossbar, in which the weights of that layer are stored in the charge or conductance state of the memory devices at the crosspoints.\n",
    "\n",
    "<center><img src=\"img/analog_Dnn.png\" style=\"width:60%; height:60%\"/></center> \n",
    "\n",
    "\n",
    "\n",
    "## IBM Analog Hardware Acceleration Kit (AIHWKIT)\n",
    "\n",
    "The IBM Analog Hardware Acceleration Kit (AIHWKIT) is an open source Python toolkit for exploring and using the capabilities of in-memory computing devices such as PCM 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",
    "\n",
    "The github repository can be found at: https://github.com/IBM/aihwkit\n",
    "\n",
    "To learn more about Analog AI and the harware befind it, refer to this webpage: https://analog-ai-demo.mybluemix.net/hardware\n",
    "\n",
    "\n",
    "### Installing the AIHWKIT\n",
    "The first thing to do is to install the AIHHKIT 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": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "29ZMdvhfzHeV",
    "outputId": "00752ae3-8164-4f2d-8e75-8da2fc2dbeb1"
   },
   "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": {
    "id": "wC9XIRtjzHeW"
   },
   "source": [
    "\n",
    "## LeNet5 Neural Network Examples\n",
    "\n",
    "In this notebook we will use the AIHWKIT to train a LeNet5 inspired analog network, as studied in the paper: https://www.frontiersin.org/articles/10.3389/fnins.2017.00538/full\n",
    "\n",
    "<img src=\"img/LeNet5_animation.png\" style=\"width:40%; height:40%\"/> \n",
    "\n",
    "The architecture of the LeNet5 network is shown below:\n",
    "\n",
    "<img src=\"img/LeNet.png\" style=\"width:40%; height:40%\"/> \n",
    "\n",
    "The network will be trained using the MNIST dataset, a collection of images representing the digits 0 to 9. \n",
    "\n",
    "From Kaggle: \"MNIST (\"Modified National Institute of Standards and Technology\") is the de facto “hello world” dataset of computer vision. Since its release in 1999, this classic dataset of handwritten images has served as the basis for benchmarking classification algorithms. As new machine learning techniques emerge, MNIST remains a reliable resource for researchers and learners alike.\"\n",
    "\n",
    "[Read more.](https://www.kaggle.com/c/digit-recognizer)\n",
    "\n",
    "<img src=\"img/MnistExamples.png\" style=\"width:40%; height:40%\"/> \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PBr6Z5FHzHeX"
   },
   "source": [
    "## Analog layers\n",
    "If the library was installed correctly, you can use the following snippet for creating an analog layer and predicting the output. In the code snippet below AnalogLinea is the Analog equivalent to Linear PyTorch layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "8-uJPmFMzHeX",
    "outputId": "d9c997cf-cd3c-47ab-c005-e474d1f98fd6"
   },
   "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": {
    "id": "jehSr26TzHeY"
   },
   "source": [
    "## RPU Configuration\n",
    "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 an rpu_config input which defines various settings of the RPU tile or the Analog hardware.\n",
    "\n",
    "Through the rpu_config parameters, the user can specify many of the hardware specs such as: device used in the cross-point array, bits used by the ADC/DAC converters, noise values and many other parameters. \n",
    "\n",
    "Additional details on the RPU configuration can be found at https://aihwkit.readthedocs.io/en/latest/using_simulator.html#rpu-configurations\n",
    "\n",
    "For this particular use case, we will define two RPU configurations that we will use later in the code. The first rpu_config uses an ideal device which is linear and symmetric in the conductance changes. The second rpu_config will use a realistic Resistive Random Access Memory (ReRAM) device with its non-idealities. We will use this two configurations to highlight their impacts on network accuracy.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using RPU configuration of an ideal Analog device\n",
    "Analog devices, when employed to implement synaptic weights of a neural network, need to meet certain specifications for the network performance to be comparable to that of a floating-point software implementation. Such specifications applicable to training a neural network on the MNIST classification benchmark were derived in [Gokmen & Vlasov, Front. Neurosci. 2016](https://www.frontiersin.org/articles/10.3389/fnins.2016.00333/full). The specifications may vary for different networks, datasets, and optimizers, and the AIHWKIT provides the capability to experiment with them. \n",
    "\n",
    "The ideal analog device is represented in the code snippet below using SingleRPUConfig(device=ConstantStepDevice()).\n",
    "This RPU configuration simulates a fictitious analog device array with idealized specifications, inspired by those listed in [Gokmen & Vlasov, Front. Neurosci. 2016](https://www.frontiersin.org/articles/10.3389/fnins.2016.00333/full). This device has an ideal linear increase of conductance with a number of pulses. It includes device-to-device variations as well as pulse-to-pulse fluctuations of conductance. The global asymmetry between positive and negative update and device-to-device asymmetry terms are set to zero. The number of steps was increased by ten-fold (to 10000 states) compared with that of Gokmen & Vlasov. \n",
    "\n",
    "<img src=\"img/idealizedDevice.png\" style=\"width:40%; height:40%\"/> "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CGNLDA2IzHeZ"
   },
   "outputs": [],
   "source": [
    "def create_rpu_config_1():\n",
    "\n",
    "    from aihwkit.simulator.configs import SingleRPUConfig\n",
    "    from aihwkit.simulator.configs.devices import ConstantStepDevice\n",
    "\n",
    "    rpu_config=SingleRPUConfig(device=ConstantStepDevice())\n",
    "\n",
    "    return rpu_config"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using RPU configuration of a realistic ReRAM device\n",
    "Resistive random-access memory (ReRAM) is a non-volatile memory technology with tuneable conductance states that can be used for in-memory computing. The conductance change of a ReRAM device is bidirectional, that is, it is possible to both increase and decrease its conductance incrementally by applying suitable electrical pulses. This capability can be exploited to implement the backpropagation algorithm. The change of conductance in oxide ReRAM is attributed to change in the configuration of the current conducting filament which consists of oxygen vacancies in a metal oxide film. \n",
    "\n",
    "The simulated ReRAM device is based on the of [Gong et al](https://www.nature.com/articles/s41467-018-04485-1). This device was fabricated with hafnium oxide as the metal oxide switching layer. The preset captures the experimentally measured response of this device to 1000 positive and 1000 negative pulses (shown in Figure 3a), including the pulse-to-pulse fluctuations. The movement of the oxygen vacancies in response to electrical signals has a probabilistic nature and it emerges as inherent randomness in conductance changes. Realistic device-to-device variability is also included in the preset to appropriately simulate the behavior of an array of such devices. \n",
    "\n",
    "<img src=\"img/reram.png\" style=\"width:40%; height:40%\"/> "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_rpu_config_2():\n",
    "\n",
    "    from aihwkit.simulator.presets import ReRamSBPreset\n",
    "\n",
    "    rpu_config=ReRamSBPreset()\n",
    "\n",
    "    return rpu_config"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "soEwkJg3zHeZ"
   },
   "source": [
    "We can now use the defined rpu_config as input of the network model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oBu6VxsJzHea"
   },
   "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": {
    "id": "EW8NXVY1zHeb"
   },
   "source": [
    "## Analog Optimizer \n",
    "\n",
    "We will use the cross entropy criteria to calculate the loss and the Stochastic Gradient Descent (SGD) as optimizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HXd7VR9szHeb"
   },
   "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": {
    "id": "iT_5bGZQzHec"
   },
   "source": [
    "## Training the network\n",
    "\n",
    "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": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "cHicu1qWzHec",
    "outputId": "e3cbef03-502c-46c0-f165-cae2ea9c6b62"
   },
   "outputs": [],
   "source": [
    "from torch import device\n",
    "from aihwkit.simulator.rpu_base import cuda\n",
    "\n",
    "\n",
    "DEVICE = device('cuda' if cuda.is_compiled() 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": {
    "id": "VwMUHVEpzHed"
   },
   "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": {
    "id": "i4-gBviwzHed"
   },
   "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": {
    "id": "8JWPjc0BzHee"
   },
   "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": {
    "id": "EU28l53ozHef"
   },
   "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",
    "    Returns:\n",
    "        nn.Module, Optimizer, Tuple: model, optimizer, and a tuple of\n",
    "            lists of train losses, validation losses, and test error\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": "markdown",
   "metadata": {
    "id": "odgAhMk2zHef"
   },
   "source": [
    "We will now download the MNIST dataset and prepare the images for the training and test:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "k_uhR68azHef"
   },
   "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": {
    "id": "qk27OXgCzHeg"
   },
   "source": [
    "### Training using the idealized rpu configuration\n",
    "Put together all the code above to train the network on the idealized analog device. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 416
    },
    "id": "b5OJ0uIbzHeg",
    "outputId": "f7109062-1a9f-424c-bfce-e38ca894d227"
   },
   "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_1()\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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training using the Analog ReRam device\n",
    "As shown by the code above Analog AI is capable of achieving high accuracy when using standard optimizer and algorithm with ideal devices. Now let's see what happens if standard SGD with the BP algorithm is used with a non-ideal device such as ReRAM"
   ]
  },
  {
   "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_2()\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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case the same network configuration with same parameters is perfoming much worse than when using the ideal device, which underscores the importance of innovating not only at a device/circuit level but also at algorithmic level.\n",
    "\n",
    "In the [next notebook](https://github.com/IBM-AI-Hardware-Center/aihwkit-notebooks/blob/main/analog_training_LeNet5_TT.ipynb) we will show how the Tiki-Taka algorithm, specifically designed for non-volatile memory and analog computing, is capable of achieving high performance also with non-ideal devices."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "Copy of LeNet5.ipynb",
   "provenance": []
  },
  "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
}