TFG2-Gpu_1024_25.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bb636ef6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip uninstall powerapi\n",
    "#!pip install powerapi\n",
    "#!pip install torchvision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f13e29e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import snntorch as snn\n",
    "from snntorch import spikeplot as splt\n",
    "from snntorch import spikegen\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import time\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0702785c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Leaky neuron model, overriding the backward pass with a custom function\n",
    "class LeakySurrogate(nn.Module):\n",
    "  def __init__(self, beta, threshold=1.0):\n",
    "      super(LeakySurrogate, self).__init__()\n",
    "\n",
    "      # initialize decay rate beta and threshold\n",
    "      self.beta = beta\n",
    "      self.threshold = threshold\n",
    "      self.spike_gradient = self.ATan.apply\n",
    "  \n",
    "  # the forward function is called each time we call Leaky\n",
    "  def forward(self, input_, mem):\n",
    "    spk = self.spike_gradient((mem-self.threshold))  # call the Heaviside function\n",
    "    reset = (self.beta * spk * self.threshold).detach() # remove reset from computational graph\n",
    "    mem = self.beta * mem + input_ - reset # Eq (1)\n",
    "    return spk, mem\n",
    "\n",
    "  # Forward pass: Heaviside function\n",
    "  # Backward pass: Override Dirac Delta with the ArcTan function\n",
    "  @staticmethod\n",
    "  class ATan(torch.autograd.Function):\n",
    "      @staticmethod\n",
    "      def forward(ctx, mem):\n",
    "          spk = (mem > 0).float() # Heaviside on the forward pass: Eq(2)\n",
    "          ctx.save_for_backward(mem)  # store the membrane for use in the backward pass\n",
    "          return spk\n",
    "\n",
    "      @staticmethod\n",
    "      def backward(ctx, grad_output):\n",
    "          (mem,) = ctx.saved_tensors  # retrieve the membrane potential \n",
    "          grad = 1 / (1 + (np.pi * mem).pow_(2)) * grad_output # Eqn 5\n",
    "          return grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7a08a0b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "lif1 = LeakySurrogate(beta=0.9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9439e055",
   "metadata": {},
   "outputs": [],
   "source": [
    "lif1 = snn.Leaky(beta=0.9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7bed8abd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataloader arguments\n",
    "batch_size = 1024\n",
    "data_path='../datos'\n",
    "\n",
    "dtype = torch.float\n",
    "dispositivo = \"cuda\"\n",
    "\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "#device = torch.device(\"cpu\")\n",
    "#device = torch_ipu.IPUDevice()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e172902b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a transform\n",
    "transform = transforms.Compose([\n",
    "            transforms.Resize((28, 28)),\n",
    "            transforms.Grayscale(),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize((0,), (1,))])\n",
    "\n",
    "mnist_train = datasets.MNIST(data_path, train=True, download=True, transform=transform)\n",
    "mnist_test = datasets.MNIST(data_path, train=False, download=True, transform=transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8f459bf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create DataLoaders\n",
    "train_loader = DataLoader(mnist_train, batch_size=batch_size, shuffle=True, drop_last=True)\n",
    "test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=True, drop_last=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3b098564",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parámetros de la red\n",
    "num_inputs = 28 * 28\n",
    "num_hidden1 = 256\n",
    "num_hidden2 = 64\n",
    "num_hidden3 = 128\n",
    "num_outputs = 10\n",
    "\n",
    "# Parámetros temporales\n",
    "num_steps = 25\n",
    "beta = 0.9\n",
    "\n",
    "# Definición de una arquitectura muy distinta\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        # Inicialización de capas\n",
    "        self.fc1 = nn.Linear(num_inputs, num_hidden1)\n",
    "        self.lif1 = snn.Leaky(beta=beta)\n",
    "        self.fc2 = nn.Linear(num_hidden1, num_hidden2)\n",
    "        self.lif2 = snn.Leaky(beta=beta)\n",
    "        self.fc3 = nn.Linear(num_hidden2, num_hidden3)\n",
    "        self.lif3 = snn.Leaky(beta=beta)\n",
    "        self.fc4 = nn.Linear(num_hidden3, num_outputs)\n",
    "        self.lif4 = snn.Leaky(beta=beta)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Inicialización de estados ocultos en t=0\n",
    "        mem1 = self.lif1.init_leaky()\n",
    "        mem2 = self.lif2.init_leaky()\n",
    "        mem3 = self.lif3.init_leaky()\n",
    "        mem4 = self.lif4.init_leaky()\n",
    "\n",
    "        # Registro de actividad de disparo\n",
    "        spk4_rec = []\n",
    "        mem4_rec = []\n",
    "\n",
    "        for step in range(num_steps):\n",
    "            cur1 = self.fc1(x)\n",
    "            spk1, mem1 = self.lif1(cur1, mem1)\n",
    "            cur2 = self.fc2(spk1)\n",
    "            spk2, mem2 = self.lif2(cur2, mem2)\n",
    "            cur3 = self.fc3(spk2)\n",
    "            spk3, mem3 = self.lif3(cur3, mem3)\n",
    "            cur4 = self.fc4(spk3)\n",
    "            spk4, mem4 = self.lif4(cur4, mem4)\n",
    "            spk4_rec.append(spk4)\n",
    "            mem4_rec.append(mem4)\n",
    "\n",
    "        return torch.stack(spk4_rec, dim=0), torch.stack(mem4_rec, dim=0)\n",
    "\n",
    "        \n",
    "# Load the network onto CUDA if available\n",
    "net = Net().to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7e40a115",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pass data into the network, sum the spikes over time\n",
    "# and compare the neuron with the highest number of spikes\n",
    "# with the target\n",
    "\n",
    "def print_batch_accuracy(data, targets, train=False):\n",
    "    output, _ = net(data.view(batch_size, -1))\n",
    "    _, idx = output.sum(dim=0).max(1)\n",
    "    acc = np.mean((targets == idx).detach().cpu().numpy())\n",
    "\n",
    "def train_printer(\n",
    "    data, targets, epoch,\n",
    "    counter, iter_counter,\n",
    "        loss_hist, test_loss_hist, test_data, test_targets):\n",
    "    #print(f\"Epoch {epoch}, Iteration {iter_counter}\")\n",
    "    #print(f\"Train Set Loss: {loss_hist[counter]:.2f}\")\n",
    "    #print(f\"Test Set Loss: {test_loss_hist[counter]:.2f}\")\n",
    "    #print_batch_accuracy(data, targets, train=True)\n",
    "    #print_batch_accuracy(test_data, test_targets, train=False)\n",
    "    print(\"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e2ba337c",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d193116d",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(net.parameters(), lr=5e-4, betas=(0.9, 0.999))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "35bc1e3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "data, targets = next(iter(train_loader))\n",
    "data = data.to(device)\n",
    "targets = targets.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "23641ac7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([25, 1024, 10])\n"
     ]
    }
   ],
   "source": [
    "spk_rec, mem_rec = net(data.view(batch_size, -1))\n",
    "#print(mem_rec.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "96f6f0ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training loss: 58.938\n"
     ]
    }
   ],
   "source": [
    "# initialize the total loss value\n",
    "loss_val = torch.zeros((1), dtype=dtype, device=device)\n",
    "\n",
    "# sum loss at every step\n",
    "for step in range(num_steps):\n",
    "  loss_val += loss(mem_rec[step], targets)\n",
    "\n",
    "#print(f\"Training loss: {loss_val.item():.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "987d83e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train set accuracy for a single minibatch: 12.40%\n"
     ]
    }
   ],
   "source": [
    "#print_batch_accuracy(data, targets, train=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "31f4b7d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# clear previously stored gradients\n",
    "optimizer.zero_grad()\n",
    "\n",
    "# calculate the gradients\n",
    "loss_val.backward()\n",
    "\n",
    "# weight update\n",
    "optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "47b9b467",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training loss: 54.836\n",
      "Train set accuracy for a single minibatch: 18.85%\n"
     ]
    }
   ],
   "source": [
    "# calculate new network outputs using the same data\n",
    "spk_rec, mem_rec = net(data.view(batch_size, -1))\n",
    "\n",
    "# initialize the total loss value\n",
    "loss_val = torch.zeros((1), dtype=dtype, device=device)\n",
    "\n",
    "# sum loss at every step\n",
    "for step in range(num_steps):\n",
    "  loss_val += loss(mem_rec[step], targets)\n",
    "\n",
    "#print(f\"Training loss: {loss_val.item():.3f}\")\n",
    "#print_batch_accuracy(data, targets, train=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "f1a273dc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0, Iteration 0\n",
      "Train Set Loss: 56.03\n",
      "Test Set Loss: 53.13\n",
      "Train set accuracy for a single minibatch: 31.84%\n",
      "Test set accuracy for a single minibatch: 28.61%\n",
      "\n",
      "\n",
      "Epoch 0, Iteration 50\n",
      "Train Set Loss: 14.94\n",
      "Test Set Loss: 13.88\n",
      "Train set accuracy for a single minibatch: 88.18%\n",
      "Test set accuracy for a single minibatch: 88.28%\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "num_epochs = 1\n",
    "loss_hist = []\n",
    "test_loss_hist = []\n",
    "counter = 0\n",
    "\n",
    "# Empezamos a medir el tiempo de entrenamiento\n",
    "start_train_time = time.time()\n",
    "\n",
    "# Outer training loop\n",
    "for epoch in range(num_epochs):\n",
    "    \n",
    "    iter_counter = 0\n",
    "    train_batch = iter(train_loader)\n",
    "\n",
    "    # Minibatch training loop\n",
    "    for data, targets in train_batch:\n",
    "        \n",
    "        data = data.to(device)\n",
    "        targets = targets.to(device)\n",
    "        \n",
    "        # forward pass\n",
    "        net.train()\n",
    "        spk_rec, mem_rec = net(data.view(batch_size, -1))\n",
    "\n",
    "        # initialize the loss & sum over time\n",
    "        loss_val = torch.zeros((1), dtype=dtype, device=device)\n",
    "        for step in range(num_steps):\n",
    "            loss_val += loss(mem_rec[step], targets)\n",
    "\n",
    "        # Gradient calculation + weight update\n",
    "        optimizer.zero_grad()\n",
    "        loss_val.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        # Store loss history for future plotting\n",
    "        loss_hist.append(loss_val.item())\n",
    "\n",
    "        # Test set\n",
    "        with torch.no_grad():\n",
    "            net.eval()\n",
    "            test_data, test_targets = next(iter(test_loader))\n",
    "            test_data = test_data.to(device)\n",
    "            test_targets = test_targets.to(device)\n",
    "\n",
    "            # Test set forward pass\n",
    "            test_spk, test_mem = net(test_data.view(batch_size, -1))\n",
    "\n",
    "            # Test set loss\n",
    "            test_loss = torch.zeros((1), dtype=dtype, device=device)\n",
    "            for step in range(num_steps):\n",
    "                test_loss += loss(test_mem[step], test_targets)\n",
    "            test_loss_hist.append(test_loss.item())\n",
    "\n",
    "            # Print train/test loss/accuracy\n",
    "            if counter % 50 == 0:\n",
    "                train_printer(\n",
    "                    data, targets, epoch,\n",
    "                    counter, iter_counter,\n",
    "                    loss_hist, test_loss_hist,\n",
    "                    test_data, test_targets)\n",
    "            counter += 1\n",
    "            iter_counter +=1\n",
    "            \n",
    "end_train_time = time.time()\n",
    "total_train_time = end_train_time - start_train_time\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "dba6cd66",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cpu\n",
      "Total correctly classified test set images: 8908/10000\n",
      "Test Set Accuracy: 89.08%\n",
      "Tiempo total: 1.8388621807098389 segundos\n",
      "Tiempo total de entrenamiento: 23.53195571899414 segundos\n"
     ]
    }
   ],
   "source": [
    "total = 0\n",
    "correct = 0\n",
    "\n",
    "# Empezamos a medir el tiempo\n",
    "start_time = time.time()\n",
    "print(f\"{dispositivo} 1024 25\")\n",
    "\n",
    "# drop_last switched to False to keep all samples\n",
    "test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=True, drop_last=False)\n",
    "# Guardar el modelo entrenado\n",
    "torch.save(net.state_dict(), '2modelo_entrenado_gpu_1024_25.pth')\n",
    "\n",
    "with torch.no_grad():\n",
    "  net.eval()\n",
    "  for data, targets in test_loader:\n",
    "\n",
    "    data = data.to(device)\n",
    "    targets = targets.to(device)\n",
    "    \n",
    "    # forward pass\n",
    "    test_spk, _ = net(data.view(data.size(0), -1))\n",
    "\n",
    "    # calculate total accuracy\n",
    "    _, predicted = test_spk.sum(dim=0).max(1)\n",
    "    total += targets.size(0)\n",
    "    correct += (predicted == targets).sum().item()\n",
    "\n",
    "end_time = time.time()\n",
    "total_time = end_time - start_time\n",
    "     \n",
    "print(f\"Total correctly classified test set images: {correct}/{total}\")\n",
    "print(f\"Test Set Accuracy: {100 * correct / total:.2f}%\")\n",
    "\n",
    "print(f\"Tiempo total: {total_time} segundos\")\n",
    "print(f\"Tiempo total de entrenamiento: {total_train_time} segundos\")"
   ]
  }
 ],
 "metadata": {
  "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}