TFG1-Gpu_256_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",
    "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 = 256\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 = 128\n",
    "num_hidden3 = 64  # Agregamos una tercera capa oculta diferente\n",
    "num_outputs = 10\n",
    "\n",
    "# Parámetros temporales\n",
    "num_steps = 25\n",
    "beta = 0.9  # Cambiamos el valor de beta\n",
    "\n",
    "# Definición de la red modificada\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)  # Agregamos una tercera capa oculta\n",
    "        self.lif3 = snn.Leaky(beta=beta)\n",
    "        self.fc4 = nn.Linear(num_hidden3, num_outputs)  # Agregamos una capa final diferente\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",
    "\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": null,
   "id": "e2ba337c",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d193116d",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(net.parameters(), lr=5e-4, betas=(0.9, 0.999))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "id": "23641ac7",
   "metadata": {},
   "outputs": [],
   "source": [
    "spk_rec, mem_rec = net(data.view(batch_size, -1))\n",
    "#print(mem_rec.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96f6f0ce",
   "metadata": {},
   "outputs": [],
   "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": null,
   "id": "987d83e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#print_batch_accuracy(data, targets, train=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "id": "47b9b467",
   "metadata": {},
   "outputs": [],
   "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": null,
   "id": "f1a273dc",
   "metadata": {},
   "outputs": [],
   "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": null,
   "id": "dba6cd66",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "total = 0\n",
    "correct = 0\n",
    "\n",
    "# Empezamos a medir el tiempo\n",
    "start_time = time.time()\n",
    "print(f\"{dispositivo} 256 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(), '1modelo_entrenado_gpu_256_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
}