NN - KAN and CNN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyPq1l8IigkoOMWZdmVNOKlc",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/TAUforPython/machinelearning/blob/main/NN%20-%20KAN%20and%20CNN.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torchvision import datasets, transforms\n",
        "from torch.utils.data import DataLoader\n",
        "import math\n",
        "import matplotlib.pyplot as plt\n",
        "from torchsummary import summary\n"
      ],
      "metadata": {
        "id": "Wg6jGk1yNj6l"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Class NN with KAN"
      ],
      "metadata": {
        "id": "wG47GwcZhs2j"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# KANLinear definition Soure: https://github.com/Blealtan/efficient-kan/blob/f39e5146af34299ad3a581d2106eb667ba0fa6fa/src/efficient_kan/kan.py#L6\n",
        "class KANLinear(torch.nn.Module):\n",
        "    def __init__(\n",
        "        self,\n",
        "        in_features,\n",
        "        out_features,\n",
        "        grid_size=5,\n",
        "        spline_order=3,\n",
        "        scale_noise=0.1,\n",
        "        scale_base=1.0,\n",
        "        scale_spline=1.0,\n",
        "        enable_standalone_scale_spline=True,\n",
        "        base_activation=torch.nn.SiLU,\n",
        "        grid_eps=0.02,\n",
        "        grid_range=[-1, 1],\n",
        "    ):\n",
        "        super(KANLinear, self).__init__()\n",
        "        self.in_features = in_features\n",
        "        self.out_features = out_features\n",
        "        self.grid_size = grid_size\n",
        "        self.spline_order = spline_order\n",
        "\n",
        "        h = (grid_range[1] - grid_range[0]) / grid_size\n",
        "        grid = (\n",
        "            (\n",
        "                torch.arange(-spline_order, grid_size + spline_order + 1) * h\n",
        "                + grid_range[0]\n",
        "            )\n",
        "            .expand(in_features, -1)\n",
        "            .contiguous()\n",
        "        )\n",
        "        self.register_buffer(\"grid\", grid)\n",
        "\n",
        "        self.base_weight = torch.nn.Parameter(torch.Tensor(out_features, in_features))\n",
        "        self.spline_weight = torch.nn.Parameter(\n",
        "            torch.Tensor(out_features, in_features, grid_size + spline_order)\n",
        "        )\n",
        "        if enable_standalone_scale_spline:\n",
        "            self.spline_scaler = torch.nn.Parameter(\n",
        "                torch.Tensor(out_features, in_features)\n",
        "            )\n",
        "\n",
        "        self.scale_noise = scale_noise\n",
        "        self.scale_base = scale_base\n",
        "        self.scale_spline = scale_spline\n",
        "        self.enable_standalone_scale_spline = enable_standalone_scale_spline\n",
        "        self.base_activation = base_activation()\n",
        "        self.grid_eps = grid_eps\n",
        "\n",
        "        self.reset_parameters()\n",
        "\n",
        "    def reset_parameters(self):\n",
        "        torch.nn.init.kaiming_uniform_(self.base_weight, a=math.sqrt(5) * self.scale_base)\n",
        "        with torch.no_grad():\n",
        "            noise = (\n",
        "                (\n",
        "                    torch.rand(self.grid_size + 1, self.in_features, self.out_features)\n",
        "                    - 1 / 2\n",
        "                )\n",
        "                * self.scale_noise\n",
        "                / self.grid_size\n",
        "            )\n",
        "            self.spline_weight.data.copy_(\n",
        "                (self.scale_spline if not self.enable_standalone_scale_spline else 1.0)\n",
        "                * self.curve2coeff(\n",
        "                    self.grid.T[self.spline_order : -self.spline_order],\n",
        "                    noise,\n",
        "                )\n",
        "            )\n",
        "            if self.enable_standalone_scale_spline:\n",
        "                # torch.nn.init.constant_(self.spline_scaler, self.scale_spline)\n",
        "                torch.nn.init.kaiming_uniform_(self.spline_scaler, a=math.sqrt(5) * self.scale_spline)\n",
        "\n",
        "    def b_splines(self, x: torch.Tensor):\n",
        "        \"\"\"\n",
        "        Compute the B-spline bases for the given input tensor.\n",
        "\n",
        "        Args:\n",
        "            x (torch.Tensor): Input tensor of shape (batch_size, in_features).\n",
        "\n",
        "        Returns:\n",
        "            torch.Tensor: B-spline bases tensor of shape (batch_size, in_features, grid_size + spline_order).\n",
        "        \"\"\"\n",
        "        assert x.dim() == 2 and x.size(1) == self.in_features\n",
        "\n",
        "        grid: torch.Tensor = (\n",
        "            self.grid\n",
        "        )  # (in_features, grid_size + 2 * spline_order + 1)\n",
        "        x = x.unsqueeze(-1)\n",
        "        bases = ((x >= grid[:, :-1]) & (x < grid[:, 1:])).to(x.dtype)\n",
        "        for k in range(1, self.spline_order + 1):\n",
        "            bases = (\n",
        "                (x - grid[:, : -(k + 1)])\n",
        "                / (grid[:, k:-1] - grid[:, : -(k + 1)])\n",
        "                * bases[:, :, :-1]\n",
        "            ) + (\n",
        "                (grid[:, k + 1 :] - x)\n",
        "                / (grid[:, k + 1 :] - grid[:, 1:(-k)])\n",
        "                * bases[:, :, 1:]\n",
        "            )\n",
        "\n",
        "        assert bases.size() == (\n",
        "            x.size(0),\n",
        "            self.in_features,\n",
        "            self.grid_size + self.spline_order,\n",
        "        )\n",
        "        return bases.contiguous()\n",
        "\n",
        "    def curve2coeff(self, x: torch.Tensor, y: torch.Tensor):\n",
        "        \"\"\"\n",
        "        Compute the coefficients of the curve that interpolates the given points.\n",
        "\n",
        "        Args:\n",
        "            x (torch.Tensor): Input tensor of shape (batch_size, in_features).\n",
        "            y (torch.Tensor): Output tensor of shape (batch_size, in_features, out_features).\n",
        "\n",
        "        Returns:\n",
        "            torch.Tensor: Coefficients tensor of shape (out_features, in_features, grid_size + spline_order).\n",
        "        \"\"\"\n",
        "        assert x.dim() == 2 and x.size(1) == self.in_features\n",
        "        assert y.size() == (x.size(0), self.in_features, self.out_features)\n",
        "\n",
        "        A = self.b_splines(x).transpose(\n",
        "            0, 1\n",
        "        )  # (in_features, batch_size, grid_size + spline_order)\n",
        "        B = y.transpose(0, 1)  # (in_features, batch_size, out_features)\n",
        "        solution = torch.linalg.lstsq(\n",
        "            A, B\n",
        "        ).solution  # (in_features, grid_size + spline_order, out_features)\n",
        "        result = solution.permute(\n",
        "            2, 0, 1\n",
        "        )  # (out_features, in_features, grid_size + spline_order)\n",
        "\n",
        "        assert result.size() == (\n",
        "            self.out_features,\n",
        "            self.in_features,\n",
        "            self.grid_size + self.spline_order,\n",
        "        )\n",
        "        return result.contiguous()\n",
        "\n",
        "    @property\n",
        "    def scaled_spline_weight(self):\n",
        "        return self.spline_weight * (\n",
        "            self.spline_scaler.unsqueeze(-1)\n",
        "            if self.enable_standalone_scale_spline\n",
        "            else 1.0\n",
        "        )\n",
        "\n",
        "    def forward(self, x: torch.Tensor):\n",
        "        assert x.dim() == 2 and x.size(1) == self.in_features\n",
        "\n",
        "        base_output = F.linear(self.base_activation(x), self.base_weight)\n",
        "        spline_output = F.linear(\n",
        "            self.b_splines(x).view(x.size(0), -1),\n",
        "            self.scaled_spline_weight.view(self.out_features, -1),\n",
        "        )\n",
        "        return base_output + spline_output\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def update_grid(self, x: torch.Tensor, margin=0.01):\n",
        "        assert x.dim() == 2 and x.size(1) == self.in_features\n",
        "        batch = x.size(0)\n",
        "\n",
        "        splines = self.b_splines(x)  # (batch, in, coeff)\n",
        "        splines = splines.permute(1, 0, 2)  # (in, batch, coeff)\n",
        "        orig_coeff = self.scaled_spline_weight  # (out, in, coeff)\n",
        "        orig_coeff = orig_coeff.permute(1, 2, 0)  # (in, coeff, out)\n",
        "        unreduced_spline_output = torch.bmm(splines, orig_coeff)  # (in, batch, out)\n",
        "        unreduced_spline_output = unreduced_spline_output.permute(\n",
        "            1, 0, 2\n",
        "        )  # (batch, in, out)\n",
        "\n",
        "        # sort each channel individually to collect data distribution\n",
        "        x_sorted = torch.sort(x, dim=0)[0]\n",
        "        grid_adaptive = x_sorted[\n",
        "            torch.linspace(\n",
        "                0, batch - 1, self.grid_size + 1, dtype=torch.int64, device=x.device\n",
        "            )\n",
        "        ]\n",
        "\n",
        "        uniform_step = (x_sorted[-1] - x_sorted[0] + 2 * margin) / self.grid_size\n",
        "        grid_uniform = (\n",
        "            torch.arange(\n",
        "                self.grid_size + 1, dtype=torch.float32, device=x.device\n",
        "            ).unsqueeze(1)\n",
        "            * uniform_step\n",
        "            + x_sorted[0]\n",
        "            - margin\n",
        "        )\n",
        "\n",
        "        grid = self.grid_eps * grid_uniform + (1 - self.grid_eps) * grid_adaptive\n",
        "        grid = torch.concatenate(\n",
        "            [\n",
        "                grid[:1]\n",
        "                - uniform_step\n",
        "                * torch.arange(self.spline_order, 0, -1, device=x.device).unsqueeze(1),\n",
        "                grid,\n",
        "                grid[-1:]\n",
        "                + uniform_step\n",
        "                * torch.arange(1, self.spline_order + 1, device=x.device).unsqueeze(1),\n",
        "            ],\n",
        "            dim=0,\n",
        "        )\n",
        "\n",
        "        self.grid.copy_(grid.T)\n",
        "        self.spline_weight.data.copy_(self.curve2coeff(x, unreduced_spline_output))\n",
        "\n",
        "    def regularization_loss(self, regularize_activation=1.0, regularize_entropy=1.0):\n",
        "        \"\"\"\n",
        "        Compute the regularization loss.\n",
        "\n",
        "        This is a dumb simulation of the original L1 regularization as stated in the\n",
        "        paper, since the original one requires computing absolutes and entropy from the\n",
        "        expanded (batch, in_features, out_features) intermediate tensor, which is hidden\n",
        "        behind the F.linear function if we want an memory efficient implementation.\n",
        "\n",
        "        The L1 regularization is now computed as mean absolute value of the spline\n",
        "        weights. The authors implementation also includes this term in addition to the\n",
        "        sample-based regularization.\n",
        "        \"\"\"\n",
        "        l1_fake = self.spline_weight.abs().mean(-1)\n",
        "        regularization_loss_activation = l1_fake.sum()\n",
        "        p = l1_fake / regularization_loss_activation\n",
        "        regularization_loss_entropy = -torch.sum(p * p.log())\n",
        "        return (\n",
        "            regularize_activation * regularization_loss_activation\n",
        "            + regularize_entropy * regularization_loss_entropy\n",
        "        )\n",
        "\n",
        "# CNN model for CIFAR-10 with KANLinear\n",
        "class CNNKAN(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(CNNKAN, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)\n",
        "        self.pool1 = nn.MaxPool2d(2)\n",
        "        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)\n",
        "        self.pool2 = nn.MaxPool2d(2)\n",
        "        self.kan1 = KANLinear(64 * 8 * 8, 256)\n",
        "        self.kan2 = KANLinear(256, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = F.selu(self.conv1(x))\n",
        "        x = self.pool1(x)\n",
        "        x = F.selu(self.conv2(x))\n",
        "        x = self.pool2(x)\n",
        "        x = x.view(x.size(0), -1)\n",
        "        x = self.kan1(x)\n",
        "        x = self.kan2(x)\n",
        "        return x\n",
        "\n",
        "class CNN(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(CNN, self).__init__()\n",
        "        # Convolutional layers\n",
        "        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)\n",
        "        self.pool1 = nn.MaxPool2d(2, 2)\n",
        "        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)\n",
        "        self.pool2 = nn.MaxPool2d(2, 2)\n",
        "\n",
        "        # Fully connected layers\n",
        "        self.fc1 = nn.Linear(64 * 8 * 8, 256)\n",
        "        self.fc2 = nn.Linear(256, 10)  # Final output layer\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Convolutional layers\n",
        "        x = F.selu(self.conv1(x))\n",
        "        x = self.pool1(x)\n",
        "        x = F.selu(self.conv2(x))\n",
        "        x = self.pool2(x)\n",
        "\n",
        "        # Flattening the layer for the fully connected layer\n",
        "        x = x.view(x.size(0), -1)\n",
        "\n",
        "        # Fully connected layers\n",
        "        x = F.selu(self.fc1(x))\n",
        "        x = self.fc2(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "def print_parameter_details(model):\n",
        "    total_params = 0\n",
        "    for name, parameter in model.named_parameters():\n",
        "        if parameter.requires_grad:\n",
        "            params = parameter.numel()  # Number of elements in the tensor\n",
        "            total_params += params\n",
        "            print(f\"{name}: {params}\")\n",
        "    print(f\"Total trainable parameters: {total_params}\")\n",
        "\n",
        "\n",
        "def visualize_kan_parameters(kan_layer, layer_name):\n",
        "    base_weights = kan_layer.base_weight.data.cpu().numpy()\n",
        "    plt.hist(base_weights.ravel(), bins=50)\n",
        "    plt.title(f\"Distribution of Base Weights - {layer_name}\")\n",
        "    plt.xlabel(\"Weight Value\")\n",
        "    plt.ylabel(\"Frequency\")\n",
        "    plt.show()\n",
        "    if hasattr(kan_layer, 'spline_weight'):\n",
        "        spline_weights = kan_layer.spline_weight.data.cpu().numpy()\n",
        "        plt.hist(spline_weights.ravel(), bins=50)\n",
        "        plt.title(f\"Distribution of Spline Weights - {layer_name}\")\n",
        "        plt.xlabel(\"Weight Value\")\n",
        "        plt.ylabel(\"Frequency\")\n",
        "        plt.show()\n"
      ],
      "metadata": {
        "id": "wipcVOztNtWO"
      },
      "execution_count": 40,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Modeling NN with Kolmogorov-Arnold-Network"
      ],
      "metadata": {
        "id": "sQzA0Bp4hzaC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "#model = CNN().to(device)\n",
        "\n",
        "# Uncommnet this line for CNN KAN.\n",
        "model = CNNKAN().to(device)\n",
        "print(model)\n",
        "print_parameter_details(model)\n",
        "summary(model,  input_size=(3, 32, 32))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0c4s8iYiOHel",
        "outputId": "9c48732b-c951-4a2b-d600-bbdb922d2582"
      },
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CNNKAN(\n",
            "  (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "  (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
            "  (conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "  (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
            "  (kan1): KANLinear(\n",
            "    (base_activation): SiLU()\n",
            "  )\n",
            "  (kan2): KANLinear(\n",
            "    (base_activation): SiLU()\n",
            "  )\n",
            ")\n",
            "conv1.weight: 864\n",
            "conv1.bias: 32\n",
            "conv2.weight: 18432\n",
            "conv2.bias: 64\n",
            "kan1.base_weight: 1048576\n",
            "kan1.spline_weight: 8388608\n",
            "kan1.spline_scaler: 1048576\n",
            "kan2.base_weight: 2560\n",
            "kan2.spline_weight: 20480\n",
            "kan2.spline_scaler: 2560\n",
            "Total trainable parameters: 10530752\n",
            "----------------------------------------------------------------\n",
            "        Layer (type)               Output Shape         Param #\n",
            "================================================================\n",
            "            Conv2d-1           [-1, 32, 32, 32]             896\n",
            "         MaxPool2d-2           [-1, 32, 16, 16]               0\n",
            "            Conv2d-3           [-1, 64, 16, 16]          18,496\n",
            "         MaxPool2d-4             [-1, 64, 8, 8]               0\n",
            "              SiLU-5                 [-1, 4096]               0\n",
            "         KANLinear-6                  [-1, 256]               0\n",
            "              SiLU-7                  [-1, 256]               0\n",
            "         KANLinear-8                   [-1, 10]               0\n",
            "================================================================\n",
            "Total params: 19,392\n",
            "Trainable params: 19,392\n",
            "Non-trainable params: 0\n",
            "----------------------------------------------------------------\n",
            "Input size (MB): 0.01\n",
            "Forward/backward pass size (MB): 0.50\n",
            "Params size (MB): 0.07\n",
            "Estimated Total Size (MB): 0.59\n",
            "----------------------------------------------------------------\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "for name, param in model.named_parameters():\n",
        "    print(f\"{name}: {param.size()} {'requires_grad' if param.requires_grad else 'frozen'}\")\n",
        "\n",
        "# TODO: Need to explore various Optimizer and optimize the Learning Rate.\n",
        "optimizer = optim.AdamW(model.parameters(), lr=0.001, weight_decay=1e-3)\n",
        "transform = transforms.Compose([\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.243, 0.261))\n",
        "])\n",
        "\n",
        "#The CIFAR-10 dataset consists of 60000 32x32 colour images in 10 classes,\n",
        "#with 6000 images per class. There are 50000 training images and 10000 test images.\n",
        "\n",
        "train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)\n",
        "test_dataset = datasets.CIFAR10(root='./data', train=False, transform=transform)\n",
        "train_loader = DataLoader(train_dataset, batch_size=500, shuffle=True)\n",
        "test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oMDyXKyLORr-",
        "outputId": "666ac18b-5d3b-49de-c42b-f6ace680f304"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "conv1.weight: torch.Size([32, 3, 3, 3]) requires_grad\n",
            "conv1.bias: torch.Size([32]) requires_grad\n",
            "conv2.weight: torch.Size([64, 32, 3, 3]) requires_grad\n",
            "conv2.bias: torch.Size([64]) requires_grad\n",
            "fc1.weight: torch.Size([256, 4096]) requires_grad\n",
            "fc1.bias: torch.Size([256]) requires_grad\n",
            "fc2.weight: torch.Size([10, 256]) requires_grad\n",
            "fc2.bias: torch.Size([10]) requires_grad\n",
            "Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to ./data/cifar-10-python.tar.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 170498071/170498071 [00:05<00:00, 29639148.19it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting ./data/cifar-10-python.tar.gz to ./data\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "len(train_dataset)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DeIuHxqKmoZC",
        "outputId": "de0f7fe0-4b17-4fe2-e84a-e1bb8481c2be"
      },
      "execution_count": 44,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "50000"
            ]
          },
          "metadata": {},
          "execution_count": 44
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def train(model, device, train_loader, optimizer, epoch):\n",
        "    model.train()\n",
        "    for batch_idx, (data, target) in enumerate(train_loader):\n",
        "        data, target = data.to(device), target.to(device)\n",
        "        optimizer.zero_grad()\n",
        "        output = model(data)\n",
        "        loss = nn.CrossEntropyLoss()(output, target)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        if batch_idx % 10 == 0:\n",
        "            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\\tLoss: {loss.item():.6f}')\n",
        "\n",
        "def evaluate(model, device, test_loader):\n",
        "    model.eval()\n",
        "    test_loss = 0\n",
        "    correct = 0\n",
        "    with torch.no_grad():\n",
        "        for data, target in test_loader:\n",
        "            data, target = data.to(device), target.to(device)\n",
        "            output = model(data)\n",
        "            test_loss += nn.CrossEntropyLoss()(output, target).item()\n",
        "            pred = output.argmax(dim=1, keepdim=True)\n",
        "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
        "    test_loss /= len(test_loader.dataset)\n",
        "    print(f'\\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({100. * correct / len(test_loader.dataset):.0f}%)\\n')\n"
      ],
      "metadata": {
        "id": "1B0ynyucOZkz"
      },
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for epoch in range(1):\n",
        "    train(model, device, train_loader, optimizer, epoch)\n",
        "    evaluate(model, device, test_loader)\n",
        "torch.save(model.state_dict(), 'model_weights_KAN.pth')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KrIc6HsqObS_",
        "outputId": "b45eec2e-a7fb-4303-e9be-f14ce4871d20"
      },
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Train Epoch: 0 [0/50000 (0%)]\tLoss: 2.301504\n",
            "Train Epoch: 0 [5000/50000 (10%)]\tLoss: 2.304529\n",
            "Train Epoch: 0 [10000/50000 (20%)]\tLoss: 2.301426\n",
            "Train Epoch: 0 [15000/50000 (30%)]\tLoss: 2.302467\n",
            "Train Epoch: 0 [20000/50000 (40%)]\tLoss: 2.301684\n",
            "Train Epoch: 0 [25000/50000 (50%)]\tLoss: 2.303531\n",
            "Train Epoch: 0 [30000/50000 (60%)]\tLoss: 2.301704\n",
            "Train Epoch: 0 [35000/50000 (70%)]\tLoss: 2.301162\n",
            "Train Epoch: 0 [40000/50000 (80%)]\tLoss: 2.300113\n",
            "Train Epoch: 0 [45000/50000 (90%)]\tLoss: 2.301013\n",
            "\n",
            "Test set: Average loss: 0.0092, Accuracy: 890/10000 (9%)\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def encode_data(model, data_loader, log_step=10, logging=print):\n",
        "    \"\"\"Encode all images and captions loadable by `data_loader`\n",
        "    \"\"\"\n",
        "    batch_time = AverageMeter()\n",
        "    val_logger = LogCollector()\n",
        "\n",
        "    # switch to evaluate mode\n",
        "    model.val_start()\n",
        "\n",
        "    end = time.time()\n",
        "\n",
        "    # numpy array to keep all the embeddings\n",
        "    img_embs = None\n",
        "    cap_embs = None\n",
        "    for i, (images, captions, lengths, ids) in enumerate(data_loader):\n",
        "        # make sure val logger is used\n",
        "        model.logger = val_logger\n",
        "\n",
        "        if(i==0 or i > 38):\n",
        "                print('     ---the {0}th images, captions, lengths, ids--from data_loader---'.format(i))\n",
        "                #print('images shape[0] {0}  shape[1] {1} shape[2] {2} shape[3] {3}'.format(images.shape[0], images.shape[1], images.shape[2], images.shape[3]))\n",
        "                print(images)\n",
        "                print('captions shape[0] {0}  shape[1] {1}'.format(captions.shape[0], captions.shape[1]))\n",
        "                print(captions)\n",
        "                print('len(lengths) {0}'.format(len(lengths)))\n",
        "                print(lengths)\n",
        "                print('len(ids) {0}'.format(len(ids)))\n",
        "                print(ids)\n",
        "                print('')\n",
        "        # compute the embeddings\n",
        "        img_emb, cap_emb = model.forward_emb(images, captions, lengths,\n",
        "                                             volatile=True)\n",
        "        if(i > 38):\n",
        "                print('     ---the {0}th img_emb, cap_emb = model.forward_emb(images, captions, lengths, volatile=True)---'.format(i))\n",
        "                print('img_emb size(0) {0}  size(1) {1}'.format(img_emb.size(0), img_emb.size(1)))\n",
        "                print(img_emb)\n",
        "                print('cap_emb size(0) {0}  size(1) {1}'.format(cap_emb.size(0), cap_emb.size(1)))\n",
        "                print(cap_emb)\n",
        "\n",
        "        # initialize the numpy arrays given the size of the embeddings\n",
        "        if img_embs is None:\n",
        "            print('img_embs is None')\n",
        "            print('len(data_loader.dataset)={0}'.format(len(data_loader.dataset)))\n",
        "            img_embs = np.zeros((len(data_loader.dataset), img_emb.size(1)))\n",
        "            cap_embs = np.zeros((len(data_loader.dataset), cap_emb.size(1)))\n",
        "\n",
        "        # preserve the embeddings by copying from gpu and converting to numpy\n",
        "        img_embs[ids] = img_emb.data.cpu().numpy().copy()\n",
        "        cap_embs[ids] = cap_emb.data.cpu().numpy().copy()\n",
        "        if(i==0 or i > 38):\n",
        "            print('          ---img_emb[0]---')\n",
        "            print(img_emb[0])\n",
        "            print('          ---img_embs[0]---')\n",
        "            print(img_embs[0])\n",
        "            print('          ---cap_emb[0]---')\n",
        "            print(cap_emb[0])\n",
        "            print('          ---cap_embs[ids[0]]---')\n",
        "            print(cap_embs[ids[0]])\n",
        "            print('          ---img_emb[5]---')\n",
        "            print(img_emb[5])\n",
        "            print('          ---img_embs[ids[5]]---')\n",
        "            print(img_embs[ids[5]])\n",
        "            print('          ---cap_emb[6]---')\n",
        "            print(cap_emb[6])\n",
        "            print('          ---cap_embs[ids[6]]---')\n",
        "            print(cap_embs[ids[6]])\n",
        "        # measure accuracy and record loss\n",
        "        model.forward_loss(img_emb, cap_emb)\n",
        "\n",
        "        # measure elapsed time\n",
        "        batch_time.update(time.time() - end)\n",
        "        end = time.time()\n",
        "\n",
        "        if i % log_step == 0:\n",
        "            logging('Test: [{0}/{1}]\\t'\n",
        "                    '{e_log}\\t'\n",
        "                    'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\\t'\n",
        "                    .format(\n",
        "                        i, len(data_loader), batch_time=batch_time,\n",
        "                        e_log=str(model.logger)))\n",
        "        del images, captions\n",
        "\n",
        "    print (i)\n",
        "    return img_embs, cap_embs"
      ],
      "metadata": {
        "id": "QyGu3XeDo2CA"
      },
      "execution_count": 47,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class AverageMeter(object):\n",
        "    \"\"\"Computes and stores the average and current value\"\"\"\n",
        "    def __init__(self):\n",
        "        self.reset()\n",
        "\n",
        "    def reset(self):\n",
        "        self.val = 0\n",
        "        self.avg = 0\n",
        "        self.sum = 0\n",
        "        self.count = 0\n",
        "\n",
        "    def update(self, val, n=1):\n",
        "        self.val = val\n",
        "        self.sum += val * n\n",
        "        self.count += n\n",
        "        self.avg = self.sum / self.count"
      ],
      "metadata": {
        "id": "BW0z6y4iqNOV"
      },
      "execution_count": 50,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "img_embs, cap_embs = encode_data(model, test_loader)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 258
        },
        "id": "ap0Zqd08XT7M",
        "outputId": "663eaa20-6121-4164-aafe-2670b1e57be4"
      },
      "execution_count": 51,
      "outputs": [
        {
          "output_type": "error",
          "ename": "NameError",
          "evalue": "name 'LogCollector' is not defined",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-51-0e94d8cb5162>\u001b[0m in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mimg_embs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcap_embs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mencode_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtest_loader\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m<ipython-input-47-a7c11d07d027>\u001b[0m in \u001b[0;36mencode_data\u001b[0;34m(model, data_loader, log_step, logging)\u001b[0m\n\u001b[1;32m      3\u001b[0m     \"\"\"\n\u001b[1;32m      4\u001b[0m     \u001b[0mbatch_time\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mAverageMeter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m     \u001b[0mval_logger\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mLogCollector\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      6\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m     \u001b[0;31m# switch to evaluate mode\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mNameError\u001b[0m: name 'LogCollector' is not defined"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#visualize_kan_parameters(kan_layer, layer_name):\n",
        "tensor.detach().cpu().numpy()\n",
        "#kan_layer.base_weight.data.cpu().numpy()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 173
        },
        "id": "vhAAmUbARQW4",
        "outputId": "9d9cdeb5-70e4-4cd8-ed09-9b85abf18feb"
      },
      "execution_count": 54,
      "outputs": [
        {
          "output_type": "error",
          "ename": "NameError",
          "evalue": "name 'tensor' is not defined",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-54-1f226ac833e6>\u001b[0m in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m#visualize_kan_parameters(kan_layer, layer_name):\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mtensor\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdetach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m \u001b[0;31m#kan_layer.base_weight.data.cpu().numpy()\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mNameError\u001b[0m: name 'tensor' is not defined"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.save(model.state_dict(), 'model_weights_KAN.pth')"
      ],
      "metadata": {
        "id": "yT8412JiSeZz"
      },
      "execution_count": 27,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "TvxbrR6tsNqH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "loaded_model = torch.load('model_weights_KAN.pth')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xauOeAEZS0jp",
        "outputId": "a7b0b1b7-7f76-4f15-b790-86066a3a5b1b"
      },
      "execution_count": 56,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "<ipython-input-56-e8ad76014b80>:1: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
            "  loaded_model = torch.load('model_weights_KAN.pth')\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(loaded_model.keys())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IG9svSHUS7FI",
        "outputId": "888f8c6f-180f-4303-a292-b9107dee540c"
      },
      "execution_count": 57,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "odict_keys(['conv1.weight', 'conv1.bias', 'conv2.weight', 'conv2.bias', 'kan1.base_weight', 'kan1.spline_weight', 'kan1.spline_scaler', 'kan1.grid', 'kan2.base_weight', 'kan2.spline_weight', 'kan2.spline_scaler', 'kan2.grid'])\n"
          ]
        }
      ]
    }
  ]
}