{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0313eb84",
   "metadata": {},
   "source": [
    "# MNIST\n",
    "\n",
    "Since QKAN (KAN-family model) can be viewed as a universal approximator, it can directly replace the job of MLP with less number of parameters.\n",
    "\n",
    "Here we are going to demonstrate how to use QKAN to do the classification task on MNIST dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a820977e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import sys\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import DataLoader\n",
    "from tqdm import tqdm\n",
    "\n",
    "from qkan import QKAN\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2c9ba985",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CNet(nn.Module):\n",
    "    def __init__(self, device):\n",
    "        super(CNet, self).__init__()\n",
    "\n",
    "        self.device = device\n",
    "\n",
    "        self.cnn1 = nn.Conv2d(\n",
    "            in_channels=1, out_channels=3, kernel_size=5, device=device\n",
    "        )\n",
    "        self.relu = nn.ReLU()\n",
    "        self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        self.cnn2 = nn.Conv2d(\n",
    "            in_channels=3, out_channels=6, kernel_size=5, device=device\n",
    "        )\n",
    "        self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        self.cnn3 = nn.Conv2d(\n",
    "            in_channels=6, out_channels=10, kernel_size=3, device=device\n",
    "        )\n",
    "        self.maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.to(self.device)\n",
    "        x = self.cnn1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.maxpool1(x)\n",
    "        x = self.cnn2(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.maxpool2(x)\n",
    "        x = self.cnn3(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.maxpool3(x)\n",
    "        x = x.flatten(start_dim=1)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "34cab282",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load MNIST\n",
    "transform = transforms.Compose(\n",
    "    [transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))]\n",
    ")\n",
    "trainset = torchvision.datasets.MNIST(\n",
    "    root=\"./data\", train=True, download=True, transform=transform\n",
    ")\n",
    "testset = torchvision.datasets.MNIST(\n",
    "    root=\"./data\", train=False, download=True, transform=transform\n",
    ")\n",
    "trainloader = DataLoader(trainset, batch_size=1000, shuffle=True)\n",
    "testloader = DataLoader(testset, batch_size=1000, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9a4b4fc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = nn.Sequential(\n",
    "    CNet(device=device),\n",
    "    QKAN([10, 10], device=device),\n",
    ")\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "accs = []\n",
    "losses = []\n",
    "test_accs = []\n",
    "test_losses = []\n",
    "test_top3_accs = []\n",
    "\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f6bf198f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 31.79it/s, accuracy=0.506, loss=1.96, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, test Loss: 1.9495182633399963, test Accuracy: 0.5352000236511231\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 36.10it/s, accuracy=0.702, loss=1.45, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2, test Loss: 1.4130778551101684, test Accuracy: 0.7213000357151031\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 35.75it/s, accuracy=0.804, loss=1.06, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3, test Loss: 1.0548779726028443, test Accuracy: 0.7998000502586364\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 34.35it/s, accuracy=0.842, loss=0.831, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4, test Loss: 0.8044245541095734, test Accuracy: 0.8453000366687775\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 35.49it/s, accuracy=0.86, loss=0.666, lr=0.001] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 5, test Loss: 0.6300603926181794, test Accuracy: 0.8797000527381897\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 35.79it/s, accuracy=0.908, loss=0.518, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 6, test Loss: 0.5075116276741027, test Accuracy: 0.903400045633316\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 34.23it/s, accuracy=0.895, loss=0.466, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 7, test Loss: 0.42405287027359007, test Accuracy: 0.9184000372886658\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 35.56it/s, accuracy=0.916, loss=0.376, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 8, test Loss: 0.3646392196416855, test Accuracy: 0.9274000406265259\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 35.60it/s, accuracy=0.929, loss=0.353, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 9, test Loss: 0.32189558148384095, test Accuracy: 0.9332000494003296\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 34.16it/s, accuracy=0.942, loss=0.284, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 10, test Loss: 0.2885587066411972, test Accuracy: 0.9390000462532043\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 35.40it/s, accuracy=0.946, loss=0.274, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 11, test Loss: 0.265482497215271, test Accuracy: 0.9436000525951386\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 35.32it/s, accuracy=0.926, loss=0.291, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 12, test Loss: 0.24529014378786088, test Accuracy: 0.9458000421524048\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 34.31it/s, accuracy=0.947, loss=0.233, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 13, test Loss: 0.22940405011177062, test Accuracy: 0.9485000431537628\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 35.25it/s, accuracy=0.953, loss=0.215, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 14, test Loss: 0.21600559800863267, test Accuracy: 0.9513000428676606\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 35.53it/s, accuracy=0.948, loss=0.226, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 15, test Loss: 0.2052382320165634, test Accuracy: 0.952800041437149\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 34.30it/s, accuracy=0.95, loss=0.205, lr=0.001] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 16, test Loss: 0.1961878292262554, test Accuracy: 0.9533000528812409\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 35.69it/s, accuracy=0.96, loss=0.176, lr=0.001] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 17, test Loss: 0.1896478660404682, test Accuracy: 0.9549000442028046\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 36.09it/s, accuracy=0.953, loss=0.199, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 18, test Loss: 0.18260392621159555, test Accuracy: 0.9564000487327575\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 34.97it/s, accuracy=0.968, loss=0.177, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 19, test Loss: 0.17513826712965966, test Accuracy: 0.9577000498771667\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 36.06it/s, accuracy=0.954, loss=0.17, lr=0.001] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 20, test Loss: 0.16960925087332726, test Accuracy: 0.9585000455379487\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(20):\n",
    "    # Train\n",
    "    model.train()\n",
    "    with tqdm(trainloader) as pbar:\n",
    "        for i, (images, labels) in enumerate(pbar):\n",
    "            optimizer.zero_grad()\n",
    "            output = model(images)\n",
    "            loss = criterion(output, labels.to(device))\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            accuracy = (output.argmax(dim=1) == labels.to(device)).float().mean()\n",
    "            accs.append(accuracy.item())\n",
    "            losses.append(loss.item())\n",
    "            pbar.set_postfix(\n",
    "                loss=loss.item(),\n",
    "                accuracy=accuracy.item(),\n",
    "                lr=optimizer.param_groups[0][\"lr\"],\n",
    "            )\n",
    "\n",
    "    # test\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    test_accuracy = 0\n",
    "    test_top3_accuracy = 0\n",
    "    with torch.no_grad():\n",
    "        for images, labels in testloader:\n",
    "            output = model(images)\n",
    "            test_loss += criterion(output, labels.to(device)).item()\n",
    "            test_accuracy += (\n",
    "                (output.argmax(dim=1) == labels.to(device)).float().mean().item()\n",
    "            )\n",
    "            test_top3_accuracy += (\n",
    "                (output.topk(3, dim=1).indices == labels.to(device).unsqueeze(1)).any(dim=1).float().mean().item()\n",
    "            )\n",
    "    test_loss /= len(testloader)\n",
    "    test_accuracy /= len(testloader)\n",
    "    test_top3_accuracy /= len(testloader)\n",
    "\n",
    "    print(f\"Epoch {epoch + 1}, test Loss: {test_loss}, test Accuracy: {test_accuracy}\")\n",
    "    test_accs.append(test_accuracy)\n",
    "    test_losses.append(test_loss)\n",
    "    test_top3_accs.append(test_top3_accuracy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "32afa55f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Top-1 Accuracy:  0.9585000455379487\n",
      "Top-3 Accuracy:  0.9934000551700592\n",
      "CNN+QKAN Model size: 1884\n"
     ]
    }
   ],
   "source": [
    "print(\"Top-1 Accuracy: \", max(test_accs))\n",
    "print(\"Top-3 Accuracy: \", max(test_top3_accs))\n",
    "\n",
    "print(\"CNN+QKAN Model size:\", len(torch.concat([param.cpu().detach().flatten() for param in model.parameters() if param.requires_grad])))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00e434fb",
   "metadata": {},
   "source": [
    "Compare to MLP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "af423d94",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CNN(nn.Module):\n",
    "    def __init__(self, device):\n",
    "        super(CNN, self).__init__()\n",
    "        self.cnet = CNet(device)\n",
    "        self.fc1 = nn.Linear(10, 20).to(device)\n",
    "        self.fc2 = nn.Linear(20, 20).to(device)\n",
    "        self.fc3 = nn.Linear(20, 10).to(device)\n",
    "        self.device = device\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x.to(self.device)\n",
    "        x = self.cnet(x)\n",
    "        x = self.fc1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.fc2(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "    \n",
    "model = CNN(device=device)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "accs = []\n",
    "losses = []\n",
    "test_accs = []\n",
    "test_losses = []\n",
    "test_top3_accs = []\n",
    "\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "77166a02",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 38.35it/s, accuracy=0.254, loss=2.07, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, test Loss: 2.06544942855835, test Accuracy: 0.24110001176595688\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 37.28it/s, accuracy=0.652, loss=1.04, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2, test Loss: 1.0108458638191222, test Accuracy: 0.6691000401973725\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 38.37it/s, accuracy=0.829, loss=0.552, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3, test Loss: 0.5260033041238785, test Accuracy: 0.8397000432014465\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 38.19it/s, accuracy=0.88, loss=0.436, lr=0.001] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4, test Loss: 0.39268899857997897, test Accuracy: 0.8824000537395478\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 37.21it/s, accuracy=0.9, loss=0.317, lr=0.001]  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 5, test Loss: 0.3284891590476036, test Accuracy: 0.9017000377178193\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 38.31it/s, accuracy=0.914, loss=0.309, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 6, test Loss: 0.2903182566165924, test Accuracy: 0.9122000455856323\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 38.54it/s, accuracy=0.912, loss=0.275, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 7, test Loss: 0.2596504405140877, test Accuracy: 0.9244000494480134\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 37.31it/s, accuracy=0.907, loss=0.293, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 8, test Loss: 0.24268167465925217, test Accuracy: 0.9282000482082366\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 36.97it/s, accuracy=0.93, loss=0.251, lr=0.001] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 9, test Loss: 0.22453297376632692, test Accuracy: 0.9329000413417816\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 37.49it/s, accuracy=0.93, loss=0.258, lr=0.001] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 10, test Loss: 0.2074310526251793, test Accuracy: 0.9391000390052795\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 37.43it/s, accuracy=0.945, loss=0.213, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 11, test Loss: 0.1980581246316433, test Accuracy: 0.9422000467777252\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 38.23it/s, accuracy=0.943, loss=0.217, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 12, test Loss: 0.19006126448512078, test Accuracy: 0.9431000411510467\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 37.67it/s, accuracy=0.954, loss=0.172, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 13, test Loss: 0.1777050383388996, test Accuracy: 0.9458000421524048\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 36.71it/s, accuracy=0.948, loss=0.175, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 14, test Loss: 0.17280858755111694, test Accuracy: 0.9464000403881073\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 38.01it/s, accuracy=0.946, loss=0.186, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 15, test Loss: 0.16691454201936723, test Accuracy: 0.949200040102005\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 37.97it/s, accuracy=0.952, loss=0.165, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 16, test Loss: 0.16387526392936708, test Accuracy: 0.9485000371932983\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 36.39it/s, accuracy=0.952, loss=0.164, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 17, test Loss: 0.15366466902196407, test Accuracy: 0.9516000390052796\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 37.86it/s, accuracy=0.945, loss=0.187, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 18, test Loss: 0.1512805711477995, test Accuracy: 0.9532000422477722\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 37.61it/s, accuracy=0.953, loss=0.158, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 19, test Loss: 0.15157018043100834, test Accuracy: 0.9514000475406647\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [00:01<00:00, 36.75it/s, accuracy=0.962, loss=0.127, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 20, test Loss: 0.1417049340903759, test Accuracy: 0.9557000458240509\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(20):\n",
    "    # Train\n",
    "    model.train()\n",
    "    with tqdm(trainloader) as pbar:\n",
    "        for i, (images, labels) in enumerate(pbar):\n",
    "            optimizer.zero_grad()\n",
    "            output = model(images)\n",
    "            loss = criterion(output, labels.to(device))\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            accuracy = (output.argmax(dim=1) == labels.to(device)).float().mean()\n",
    "            accs.append(accuracy.item())\n",
    "            losses.append(loss.item())\n",
    "            pbar.set_postfix(\n",
    "                loss=loss.item(),\n",
    "                accuracy=accuracy.item(),\n",
    "                lr=optimizer.param_groups[0][\"lr\"],\n",
    "            )\n",
    "\n",
    "    # test\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    test_accuracy = 0\n",
    "    test_top3_accuracy = 0\n",
    "    with torch.no_grad():\n",
    "        for images, labels in testloader:\n",
    "            output = model(images)\n",
    "            test_loss += criterion(output, labels.to(device)).item()\n",
    "            test_accuracy += (\n",
    "                (output.argmax(dim=1) == labels.to(device)).float().mean().item()\n",
    "            )\n",
    "            test_top3_accuracy += (\n",
    "                (output.topk(3, dim=1).indices == labels.to(device).unsqueeze(1)).any(dim=1).float().mean().item()\n",
    "            )\n",
    "    test_loss /= len(testloader)\n",
    "    test_accuracy /= len(testloader)\n",
    "    test_top3_accuracy /= len(testloader)\n",
    "\n",
    "    print(f\"Epoch {epoch + 1}, test Loss: {test_loss}, test Accuracy: {test_accuracy}\")\n",
    "    test_accs.append(test_accuracy)\n",
    "    test_losses.append(test_loss)\n",
    "    test_top3_accs.append(test_top3_accuracy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a574670d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Top-1 Accuracy:  0.9557000458240509\n",
      "Top-3 Accuracy:  0.995000034570694\n",
      "CNN+MLP Model size: 1934\n"
     ]
    }
   ],
   "source": [
    "print(\"Top-1 Accuracy: \", max(test_accs))\n",
    "print(\"Top-3 Accuracy: \", max(test_top3_accs))\n",
    "\n",
    "print(\"CNN+MLP Model size:\", len(torch.concat([param.cpu().detach().flatten() for param in model.parameters() if param.requires_grad])))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18b4f009",
   "metadata": {},
   "source": [
    "The results show that QKAN can achieve better performance than MLP, but with less number of parameters."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}