Tutorial 6_Surrogate Gradient Descent in a Convolutional SNN.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "collapsed": false,
        "id": "uSGZ6cdmpknm",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "[<img src='https://github.com/jeshraghian/snntorch/blob/master/docs/_static/img/snntorch_alpha_w.png?raw=true' width=\"400\">](https://github.com/jeshraghian/snntorch/)\n",
        "\n",
        "# snnTorch - Surrogate Gradient Descent in a Convolutional Spiking Neural Network\n",
        "## Tutorial 6\n",
        "### By Jason K. Eshraghian (www.ncg.ucsc.edu)\n",
        "### 偏白话文翻译： manesec\n",
        "\n",
        "<a href=\"https://colab.research.google.com/github/jeshraghian/snntorch/blob/master/examples/tutorial_6_CNN.ipynb\">\n",
        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
        "</a>\n",
        "\n",
        "[<img src='https://github.com/jeshraghian/snntorch/blob/master/docs/_static/img/GitHub-Mark-Light-120px-plus.png?raw=true' width=\"28\">](https://github.com/jeshraghian/snntorch/) [<img src='https://github.com/jeshraghian/snntorch/blob/master/docs/_static/img/GitHub_Logo_White.png?raw=true' width=\"80\">](https://github.com/jeshraghian/snntorch/)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rugeYYiqsrlc"
      },
      "source": [
        "snnTorch 教程系列基于以下论文。如果您发现这些资源或代码对您有帮助，请考虑引用以下来源：\n",
        "\n",
        "> <cite> [Jason K. Eshraghian, Max Ward, Emre Neftci, Xinxin Wang, Gregor Lenz, Girish Dwivedi, Mohammed Bennamoun, Doo Seok Jeong, and Wei D. Lu. \"Training Spiking Neural Networks Using Lessons From Deep Learning\". arXiv preprint arXiv:2109.12894, September 2021.](https://arxiv.org/abs/2109.12894) </cite>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ymi3sqJg28OQ"
      },
      "source": [
        "# 简介\n",
        "在本教程中，您将\n",
        "* 学习如何修改代理梯度下降算法以克服死神经元问题\n",
        "* 构建并训练一个卷积尖峰神经网络\n",
        "* 使用顺序容器 `nn.Sequential` 简化模型构建\n",
        "\n",
        "本教程的部分内容受到了 Friedemann Zenke\n",
        "关于 SNN 的大量工作的启发。请查看他关于代梯度的资料库\n",
        "[这里](https://github.com/fzenke/spytorch)，以及我最喜欢的一篇论文： E. O. Neftci、H. Mostafa、F. Zenke，[代梯度学习： 将基于梯度的优化功能引入尖峰神经网络。](https://ieeexplore.ieee.org/document/8891809) IEEE\n",
        "信号处理杂志》36，51-63。\n",
        "\n",
        "在教程的最后，我们将使用 MNIST 数据集训练一个卷积尖峰神经网络 (CSNN)，以执行图像分类。本教程的背景理论沿袭自 [教程 2、4 和 5](https://snntorch.readthedocs.io/en/latest/tutorials/index.html)，如果您需要复习，请随时查阅。\n",
        "\n",
        "如果在 Google Colab 中运行：\n",
        "* 你可以通过检查 `Runtime` > `Change runtime type` > `Hardware accelerator： GPU`\n",
        "* 接下来，点击下面的单元格并按下 `Shift+Enter` 键，安装 snnTorch 的最新 PyPi 发行版。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5tn_wUlopkon",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "!pip install snntorch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QXZ6Tuqc9Q-l"
      },
      "outputs": [],
      "source": [
        "# imports\n",
        "import snntorch as snn\n",
        "from snntorch import surrogate\n",
        "from snntorch import backprop\n",
        "from snntorch import functional as SF\n",
        "from snntorch import utils\n",
        "from snntorch import spikeplot as splt\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 torch.nn.functional as F\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import itertools"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gt2xMbLY9dVE"
      },
      "source": [
        "# 1. 替代梯度下降"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zJddJWoa0GT6"
      },
      "source": [
        "[教程 5](https://snntorch.readthedocs.io/en/latest/tutorials/index.html)提出了**死神经元的问题**，这是因为尖峰的不可微分性：\n",
        "\n",
        "$$S[t] = \\Theta(U[t] - U_{\\rm thr}) \\tag{1}$$\n",
        "$$\\frac{\\partial S}{\\partial U} = \\delta(U - U_{\\rm thr}) \\in \\{0, \\infty\\} \\tag{2}$$\n",
        "\n",
        "其中 $\\Theta(\\cdot)$ 是单位阶跃函数，$\\delta(\\cdot)$ 是狄拉克-德尔塔 (Dirac-Delta) 函数。我们之前曾使用阈值偏移 *ArcTangent* 函数的梯度来解决这个问题。\n",
        "\n",
        "其他常用的平滑函数包括 sigmoid 函数或快速 sigmoid 函数。sigmoidal 类型的函数要进行移位，以便以阈值 $U_{rm thr}$ 为中心。将膜电位的过驱动定义为 $U_{OD} = U - U_{\\rm thr}$：\n",
        "\n",
        "$$\\tilde{S} = \\frac{U_{OD}}{1+k|U_{OD}|} \\tag{3}$$\n",
        "$$\\frac{\\partial \\tilde{S}}{\\partial U} = \\frac{1}{(k|U_{OD}|+1)^2}\\tag{4}$$\n",
        "\n",
        "其中，$k$ 用来调节代替梯度函数的平滑程度，并被视为一个超参数。随着 $k$ 的增加，近似值会向 $(2)$ 中的原始导数收敛：\n",
        "\n",
        "$$\\frac{\\partial \\tilde{S}}{\\partial U} \\Bigg|_{k \\rightarrow \\infty} = \\delta(U-U_{\\rm thr})$$\n",
        "\n",
        "\n",
        "<center>\n",
        "<img src='https://github.com/jeshraghian/snntorch/blob/master/docs/_static/img/examples/tutorial6/surrogate.png?raw=true' width=\"800\">\n",
        "</center>\n",
        "\n",
        "总结一下：\n",
        "\n",
        "+ **前向传递**\n",
        "  - 确定 $(1)$ 中的单位跃迁函数的偏移量 $S$\n",
        "  - 存储 $U$ 以备后向运算时使用\n",
        "+ **后向传递**\n",
        "  - 将 $U$ 输入 $(4)$ 以计算导数项\n",
        "\n",
        "与[教程 5](https://snntorch.readthedocs.io/en/latest/tutorials/index.html)中使用的 *ArcTangent* 方法相同、快速 sigmoid 的梯度函数 可以在 泄漏整合发放（LIF）的神经元模型中超越 Dirac-Delta 函数："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5lqpuT--bZmJ"
      },
      "outputs": [],
      "source": [
        "# Leaky neuron model, overriding the backward pass with a custom function\n",
        "class LeakySigmoidSurrogate(nn.Module):\n",
        "  def __init__(self, beta, threshold=1.0, k=25):\n",
        "\n",
        "      # Leaky_Surrogate is defined in the previous tutorial and not used here\n",
        "      super(Leaky_Surrogate, self).__init__()\n",
        "\n",
        "      # initialize decay rate beta and threshold\n",
        "      self.beta = beta\n",
        "      self.threshold = threshold\n",
        "      self.surrogate_func = self.FastSigmoid.apply\n",
        "\n",
        "  # the forward function is called each time we call Leaky\n",
        "  def forward(self, input_, mem):\n",
        "    spk = self.surrogate_func((mem-self.threshold))  # call the Heaviside function\n",
        "    reset = (spk - self.threshold).detach()\n",
        "    mem = self.beta * mem + input_ - reset\n",
        "    return spk, mem\n",
        "\n",
        "  # Forward pass: Heaviside function\n",
        "  # Backward pass: Override Dirac Delta with gradient of fast sigmoid\n",
        "  @staticmethod\n",
        "  class FastSigmoid(torch.autograd.Function):\n",
        "    @staticmethod\n",
        "    def forward(ctx, mem, k=25):\n",
        "        ctx.save_for_backward(mem) # store the membrane potential for use in the backward pass\n",
        "        ctx.k = k\n",
        "        out = (mem > 0).float() # Heaviside on the forward pass: Eq(1)\n",
        "        return out\n",
        "\n",
        "    @staticmethod\n",
        "    def backward(ctx, grad_output):\n",
        "        (mem,) = ctx.saved_tensors  # retrieve membrane potential\n",
        "        grad_input = grad_output.clone()\n",
        "        grad = grad_input / (ctx.k * torch.abs(mem) + 1.0) ** 2  # gradient of fast sigmoid on backward pass: Eq(4)\n",
        "        return grad, None"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4aQvFG7ajpzU"
      },
      "source": [
        "更妙的是，所有这些都可以通过使用 snnTorch 的内置模块 `snn.surrogate` 来简化，其中来自 $(4)$ 的 $k$ 表示 `slope`。代梯度作为参数传递到 `spike_grad` 中："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2dCWD_qajyLw"
      },
      "outputs": [],
      "source": [
        "spike_grad = surrogate.fast_sigmoid(slope=25)\n",
        "beta = 0.5\n",
        "\n",
        "lif1 = snn.Leaky(beta=beta, spike_grad=spike_grad)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ffn7D6omkj5r"
      },
      "source": [
        "要探索其他可用的代梯度函数，[请查看此处的文档。](https://snntorch.readthedocs.io/en/latest/snntorch.surrogate.html)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wgzf83HE2BeB"
      },
      "source": [
        "# 2. 设置 卷积脉冲神经网络（CSNN）\n",
        "## 2.1 设置数据加载器"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pxO32dntlOB2"
      },
      "outputs": [],
      "source": [
        "# dataloader arguments\n",
        "batch_size = 128\n",
        "data_path='/data/mnist'\n",
        "\n",
        "dtype = torch.float\n",
        "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pE7eGTnulSBA"
      },
      "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)\n",
        "\n",
        "# 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": "markdown",
      "metadata": {
        "id": "6v8fBXrVlY3f"
      },
      "source": [
        "## 2.2 定义网络\n",
        "\n",
        "使用的卷积网络结构是：`12C5-MP2-64C5-MP2-1024FC10`\n",
        "\n",
        "- `12C5` 是一个  5 $\\times$ 5 的卷积核，有 12 个滤波器\n",
        "- `MP2` 是一个 2 $\\times$ 2 最大池化函数\n",
        "- `1024FC10` 是一个全连接层，将 1,024 个神经元映射到 10 个输出\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "foos_NlopDrb"
      },
      "outputs": [],
      "source": [
        "# neuron and simulation parameters\n",
        "spike_grad = surrogate.fast_sigmoid(slope=25)\n",
        "beta = 0.5\n",
        "num_steps = 50"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X4sd8PDSlGZb"
      },
      "outputs": [],
      "source": [
        "# Define Network\n",
        "class Net(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "\n",
        "        # Initialize layers\n",
        "        self.conv1 = nn.Conv2d(1, 12, 5)\n",
        "        self.lif1 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
        "        self.conv2 = nn.Conv2d(12, 64, 5)\n",
        "        self.lif2 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
        "        self.fc1 = nn.Linear(64*4*4, 10)\n",
        "        self.lif3 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
        "\n",
        "    def forward(self, x):\n",
        "\n",
        "        # Initialize hidden states and outputs at t=0\n",
        "        mem1 = self.lif1.init_leaky()\n",
        "        mem2 = self.lif2.init_leaky()\n",
        "        mem3 = self.lif3.init_leaky()\n",
        "\n",
        "        cur1 = F.max_pool2d(self.conv1(x), 2)\n",
        "        spk1, mem1 = self.lif1(cur1, mem1)\n",
        "\n",
        "        cur2 = F.max_pool2d(self.conv2(spk1), 2)\n",
        "        spk2, mem2 = self.lif2(cur2, mem2)\n",
        "\n",
        "        cur3 = self.fc1(spk2.view(batch_size, -1))\n",
        "        spk3, mem3 = self.lif3(cur3, mem3)\n",
        "        return spk3, mem3"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HVn3aYAUnWqH"
      },
      "source": [
        "在上一教程中，网络被封装在一个类中，如上图所示。\n",
        "随着网络复杂性的增加，这将增加大量的模板代码，我们可能希望避免这种情况。另外，也可以使用 `nn.Sequential` 方法。\n",
        "\n",
        "> 注意：下面的代码块模拟的是单个时间步长，需要单独的 for 循环。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AoYBY89angvp"
      },
      "outputs": [],
      "source": [
        "#  Initialize Network\n",
        "net = nn.Sequential(nn.Conv2d(1, 12, 5),\n",
        "                    nn.MaxPool2d(2),\n",
        "                    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True),\n",
        "                    nn.Conv2d(12, 64, 5),\n",
        "                    nn.MaxPool2d(2),\n",
        "                    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True),\n",
        "                    nn.Flatten(),\n",
        "                    nn.Linear(64*4*4, 10),\n",
        "                    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True)\n",
        "                    ).to(device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7Qgw1dRmpOlo"
      },
      "source": [
        "`init_hidden` 参数用来初始化隐藏状态（这里是膜电位），这可以作为实例变量以降低代码复杂性。\n",
        "\n",
        "如果激活了 `init_hidden`，膜电位就不会显式地返回给用户，从而确保只有输出尖峰会按顺序通过由 `nn.Sequential` 封装的层。\n",
        "\n",
        "要获得最后一层的输出，请设置参数`output=True`，这使得最后一层也能够返回神经元的尖峰和膜电位响应。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A-bCSQmBstvd"
      },
      "source": [
        "## 2.3 前向传递\n",
        "\n",
        "模拟持续时间`num_steps`的前向传播如下所示："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IxPPMND-pMxE"
      },
      "outputs": [],
      "source": [
        "data, targets = next(iter(train_loader))\n",
        "data = data.to(device)\n",
        "targets = targets.to(device)\n",
        "\n",
        "for step in range(num_steps):\n",
        "    spk_out, mem_out = net(data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C3PxtCobuH_e"
      },
      "source": [
        "将其包装在一个函数中，记录随时间变化的膜电位和尖峰响应："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ykdnD3tRuHcs"
      },
      "outputs": [],
      "source": [
        "def forward_pass(net, num_steps, data):\n",
        "  mem_rec = []\n",
        "  spk_rec = []\n",
        "  utils.reset(net)  # resets hidden states for all LIF neurons in net\n",
        "\n",
        "  for step in range(num_steps):\n",
        "      spk_out, mem_out = net(data)\n",
        "      spk_rec.append(spk_out)\n",
        "      mem_rec.append(mem_out)\n",
        "\n",
        "  return torch.stack(spk_rec), torch.stack(mem_rec)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "unJrx3pXcXii"
      },
      "outputs": [],
      "source": [
        "spk_rec, mem_rec = forward_pass(net, num_steps, data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zqJdfllYbc16"
      },
      "source": [
        "# 3. 训练循环"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x-kquOWLY1Jo"
      },
      "source": [
        "## 3.1 损失函数使用 snn.Functional"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MlA56BgOYq1D"
      },
      "source": [
        "在之前的教程中，使用输出神经元的膜电位与目标之间的交叉熵损失来训练网络。\n",
        "\n",
        "这次，每个神经元的尖峰总数将用于计算交叉熵。\n",
        "\n",
        "`snn.function`模块中包含各种损失函数，类似于 PyTorch 中的 `torch.nn.function`。\n",
        "\n",
        "这些函数混合使用交叉熵和均方误差，可以应用于尖峰或膜电位，以训练 速率编码 或 延迟编码 的网络。\n",
        "\n",
        "下面的方法将交叉熵损失应用于输出尖峰计数，以便训练速率编码网络："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UZ2BM6d6a11l"
      },
      "outputs": [],
      "source": [
        "# already imported snntorch.functional as SF\n",
        "loss_fn = SF.ce_rate_loss()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q39HCIeOa4fC"
      },
      "source": [
        "尖峰的记录作为第一个参数传递给`loss_fn`，目标神经元索引作为第二个参数传递以生成损失。 [该文档提供了更多信息和示例。](https://snntorch.readthedocs.io/en/latest/snntorch.function.html#snntorch.function.ce_rate_loss)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xEVzMvujcjsE"
      },
      "outputs": [],
      "source": [
        "loss_val = loss_fn(spk_rec, targets)\n",
        "\n",
        "print(f\"The loss from an untrained network is {loss_val.item():.3f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HS_zeb5mbqjw"
      },
      "source": [
        "## 3.2 精度使用 snn.Functional\n",
        "\n",
        "`SF.accuracy_rate()` 函数是以预测输出尖峰和实际目标作为参数。`accuracy_rate` 假设使用的是速率代码来解释输出，方法是检查尖峰计数最高的神经元的索引是否与目标索引相匹配。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yq7_rly0c9b4"
      },
      "outputs": [],
      "source": [
        "acc = SF.accuracy_rate(spk_rec, targets)\n",
        "\n",
        "print(f\"The accuracy of a single batch using an untrained network is {acc*100:.3f}%\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r4Z6bnqCdL50"
      },
      "source": [
        "由于上述函数只返回单批数据的准确度，因此可以写一个函数返回整个 DataLoader 对象的精确度："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IqxDKFvrdXuF"
      },
      "outputs": [],
      "source": [
        "def batch_accuracy(train_loader, net, num_steps):\n",
        "  with torch.no_grad():\n",
        "    total = 0\n",
        "    acc = 0\n",
        "    net.eval()\n",
        "\n",
        "    train_loader = iter(train_loader)\n",
        "    for data, targets in train_loader:\n",
        "      data = data.to(device)\n",
        "      targets = targets.to(device)\n",
        "      spk_rec, _ = forward_pass(net, num_steps, data)\n",
        "\n",
        "      acc += SF.accuracy_rate(spk_rec, targets) * spk_rec.size(1)\n",
        "      total += spk_rec.size(1)\n",
        "\n",
        "  return acc/total"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_u43hKAvefWM"
      },
      "outputs": [],
      "source": [
        "test_acc = batch_accuracy(test_loader, net, num_steps)\n",
        "\n",
        "print(f\"The total accuracy on the test set is: {test_acc * 100:.2f}%\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F1pzWXlsYoIu"
      },
      "source": [
        "## 3.3 训练循环"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KAUDYl3gf0G-"
      },
      "source": [
        "以下训练循环在性质上与之前的教程相似"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "y_VQ9es-gSO3"
      },
      "outputs": [],
      "source": [
        "optimizer = torch.optim.Adam(net.parameters(), lr=1e-2, betas=(0.9, 0.999))\n",
        "num_epochs = 1\n",
        "loss_hist = []\n",
        "test_acc_hist = []\n",
        "counter = 0\n",
        "\n",
        "# Outer training loop\n",
        "for epoch in range(num_epochs):\n",
        "\n",
        "    # Training loop\n",
        "    for data, targets in iter(train_loader):\n",
        "        data = data.to(device)\n",
        "        targets = targets.to(device)\n",
        "\n",
        "        # forward pass\n",
        "        net.train()\n",
        "        spk_rec, _ = forward_pass(net, num_steps, data)\n",
        "\n",
        "        # initialize the loss & sum over time\n",
        "        loss_val = loss_fn(spk_rec, 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",
        "        if counter % 50 == 0:\n",
        "          with torch.no_grad():\n",
        "              net.eval()\n",
        "\n",
        "              # Test set forward pass\n",
        "              test_acc = batch_accuracy(test_loader, net, num_steps)\n",
        "              print(f\"Iteration {counter}, Test Acc: {test_acc * 100:.2f}%\\n\")\n",
        "              test_acc_hist.append(test_acc.item())\n",
        "\n",
        "        counter += 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TjRPDFWxj2eS"
      },
      "source": [
        "尽管选择了一些相当通用的值和架构，但考虑到短暂的训练运行，测试集的准确性应该具有相当的可观的！"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "collapsed": false,
        "id": "HxU7P7xFpko3",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "# 4. 结果\n",
        "## 4.1 绘图测试准确度"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_Pk_EScnpkpj",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "# Plot Loss\n",
        "fig = plt.figure(facecolor=\"w\")\n",
        "plt.plot(test_acc_hist)\n",
        "plt.title(\"Test Set Accuracy\")\n",
        "plt.xlabel(\"Epoch\")\n",
        "plt.ylabel(\"Accuracy\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nYFamUJLkVY3"
      },
      "source": [
        "## 4.2 尖峰计数器"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MDE3ms9ulo-t"
      },
      "source": [
        "对一批数据前向运行传递以获得尖峰和膜读数"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dqIjaw1kk4-6"
      },
      "outputs": [],
      "source": [
        "spk_rec, mem_rec = forward_pass(net, num_steps, data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n4QiJx2HlkMH"
      },
      "source": [
        "修改 `idx` 允许索引模拟小批量中的各种样本。 可以使用`splt.spike_count`探索几个不同样本的尖峰行为！\n",
        "\n",
        "> 注意：如果您在桌面上本地运行笔记本，请取消注释下面的行并修改 ffmpeg.exe 的路径\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4HcwxfC6kfy0"
      },
      "outputs": [],
      "source": [
        "from IPython.display import HTML\n",
        "\n",
        "idx = 0\n",
        "\n",
        "fig, ax = plt.subplots(facecolor='w', figsize=(12, 7))\n",
        "labels=['0', '1', '2', '3', '4', '5', '6', '7', '8','9']\n",
        "print(f\"The target label is: {targets[idx]}\")\n",
        "\n",
        "plt.rcParams['animation.ffmpeg_path'] = 'D:\\\\GoodApp\\\\ffmpeg\\\\bin\\\\ffmpeg.exe'\n",
        "\n",
        "#  Plot spike count histogram\n",
        "anim = splt.spike_count(spk_rec[:, idx].detach().cpu(), fig, ax, labels=labels,\n",
        "                        animate=True, interpolate=4)\n",
        "\n",
        "HTML(anim.to_html5_video())\n",
        "# anim.save(\"spike_bar.mp4\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s0dAgWUt2o6E"
      },
      "source": [
        "# 结论\n",
        "你现在应该掌握了 snnTorch 的基本功能，并能够开始运行自己的实验。 [在下一个教程中](https://snntorch.readthedocs.io/en/latest/tutorials/index.html)，我们将使用神经形态数据集训练网络。\n",
        "\n",
        "特别感谢 [Gianfrancesco Angelini](https://github.com/gianfa) 提供有关本教程的宝贵反馈。\n",
        "\n",
        "如果您喜欢这个项目，请考虑为 GitHub 上的存储库加⭐，因为这是支持它的最简单、最好的方式。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q8eFeW-yeGci"
      },
      "source": [
        "# Additional Resources\n",
        "* [Check out the snnTorch GitHub project here.](https://github.com/jeshraghian/snntorch)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "tutorial_6_CNN.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.10"
    },
    "vscode": {
      "interpreter": {
        "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}