HSi_UH2013_P7_AB_VIM_V3_6_0330.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Judyxyang/judyxyang/blob/master/HSi_UH2013_P7_AB_VIM_V3_6_0330.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7adnkHUo9TBz"
      },
      "source": [
        "# HyperMamba Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "a3DW2ucDRG2Z",
        "outputId": "0a8f5600-d439-4fb7-efcd-9e1aa8e8a0f0"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting spectral\n",
            "  Downloading spectral-0.23.1-py3-none-any.whl (212 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m212.9/212.9 kB\u001b[0m \u001b[31m4.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting mat73\n",
            "  Downloading mat73-0.63-py3-none-any.whl (19 kB)\n",
            "Collecting einops\n",
            "  Downloading einops-0.7.0-py3-none-any.whl (44 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m44.6/44.6 kB\u001b[0m \u001b[31m4.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from spectral) (1.25.2)\n",
            "Requirement already satisfied: h5py in /usr/local/lib/python3.10/dist-packages (from mat73) (3.9.0)\n",
            "Installing collected packages: spectral, einops, mat73\n",
            "Successfully installed einops-0.7.0 mat73-0.63 spectral-0.23.1\n"
          ]
        }
      ],
      "source": [
        "pip install spectral mat73  einops"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "zudwFZMrRDDC"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import sys\n",
        "import os\n",
        "import math\n",
        "\n",
        "from einops import rearrange\n",
        "import torch.optim as optim\n",
        "import torchvision\n",
        "from torchvision import datasets, transforms\n",
        "from scipy import io\n",
        "import torch.utils.data\n",
        "import scipy.io as sio\n",
        "import mat73\n",
        "import matplotlib.pyplot as plt\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OATt9TAf90MC"
      },
      "source": [
        "# 0 Upload Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Vpwq4Yi-tgjs",
        "outputId": "0610ebac-4055-4dc2-e0c3-b24e1f220849"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vMin2GFJtorP",
        "outputId": "c63324d5-357d-48a2-c89d-4c1c6127a51f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            " 2013_DFTC\n",
            " 2013_IEEE_GRSS_DF_Contest_CASI_349_1905_144.mat\n",
            " 2013_IEEE_GRSS_DF_Contest_LiDAR.mat\n",
            " ablationp5_spatialonly_UH2013_model_state_dict.pth\n",
            " ablationp7_removeZpath_Uh2013_model_state_dict.pth\n",
            " ablationp7_removeZpath_UH2013_model_state_dict.pth\n",
            " ablationp7_spatialonly_Uh2013_model_state_dict.pth\n",
            " ablationp7_spatialonly_UH2013_model_state_dict.pth\n",
            " ablationp7_Uh2013_model_state_dict.pth\n",
            " ablationp7_UH2013_model_state_dict.pth\n",
            " ablation_Uh2013_model_state_dict.pth\n",
            " ablation_UH2013_model_state_dict.pth\n",
            " Autoencodermodel.pth\n",
            " Autoencodermodel_uh2013_20Ksample.pth\n",
            " Autoencodermodel_uh2013_adam_20Ksample.pth\n",
            " Autoencodermodel_uh2013_adamp13_20Ksample.pth\n",
            " Autoencodermodel_uh2013_adamp9_20Ksample.pth\n",
            " Autoencodermodel_uh2013_admp3fuseddata_50Ksample.pth\n",
            " Autoencodermodel_uh2013_admp5fuseddata_50Ksample.pth\n",
            " Autoencodermodel_uh2013_admp7fuseddata_50Ksample.pth\n",
            " Autoencodermodel_uh2013.pth\n",
            " Autoencodermodel_uh2013_rms_20Ksample.pth\n",
            " Autoencodermodel_uh2013_sgdp11fuseddata_20Ksample.pth\n",
            " Autoencodermodel_uh2013_sgdp13_20Ksample.pth\n",
            " Autoencodermodel_uh2013_sgdp13fuseddata_20Ksample.pth\n",
            " Autoencodermodel_uh2013_sgdp3_20Ksample.pth\n",
            " Autoencodermodel_uh2013_sgdp3fuseddata_20Ksample.pth\n",
            " Autoencodermodel_uh2013_sgdp3fuseddata_50Ksample.pth\n",
            " Autoencodermodel_uh2013_sgdp5_20Ksample.pth\n",
            " Autoencodermodel_uh2013_sgdp5fuseddata_20Ksample.pth\n",
            " Autoencodermodel_uh2013_sgdp5fuseddata_50Ksample.pth\n",
            " Autoencodermodel_uh2013_sgdp7_20Ksample.pth\n",
            " Autoencodermodel_uh2013_sgdp7fuseddata_20Ksample.pth\n",
            " Autoencodermodel_uh2013_sgdp7fuseddata_50Ksample.pth\n",
            " Autoencodermodel_uh2013_sgdp9fuseddata_20Ksample.pth\n",
            " Autoencodermodel_uh2013_sgdp9fuseddata_50Ksample.pth\n",
            " Autoencodermodel_uh2013_sgdp9v1_20Ksample.pth\n",
            " best_contrastive_model.pth\n",
            " best_hyperspectral_lidar_fusion_mamba_model.pth\n",
            " best_model_p11_weights.pth\n",
            " best_model_p13_weights.pth\n",
            " best_model_p15_weights.pth\n",
            " best_model_p3_weights.pth\n",
            " best_model_p7_weights.pth\n",
            " best_model_p9_weights.pth\n",
            " best_model.pth\n",
            " best_model_weights.pth\n",
            " BS_RLFCC_UH2013.ipynb\n",
            " cae_model_state_batchscore.pth\n",
            " cae_model_state.pth\n",
            "'Contrastive Best model'\n",
            " cross_modal_autoencoder_best.pth\n",
            " cross_modal_autoencoder.pth\n",
            " DFTC2013_Fusion_Model.h5\n",
            " GRSS2013.mat\n",
            " hyperspectral_lidar_fusion_mamba_model.pth\n",
            " __MACOSX\n",
            " model_path.pth\n",
            " model.pth\n",
            " model_state_dict.pth\n",
            " output_mask.tif\n",
            " p5_Uh2013_model_state_dict.pth\n",
            " p5_UH2013_model_state_dict.pth\n",
            " p7_Uh2013_model_state_dict.pth\n",
            " p7_UH2013_model_state_dict.pth\n",
            " saved_model.pt\n",
            " trained_model.pth\n",
            " Uh2013_model_state_dict.pth\n",
            " UH2013_model_state_dict.pth\n",
            " vis_attn_best_model.pt\n"
          ]
        }
      ],
      "source": [
        "! ls '/content/drive/MyDrive/A02_RemoteSensingData/UHS_2013_DFTC/'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "4NMT_vf3t3Bg"
      },
      "outputs": [],
      "source": [
        "# # Define the path\n",
        "path='/content/drive/MyDrive/A02_RemoteSensingData/UHS_2013_DFTC/'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3d23b73f",
        "outputId": "aa769dc9-9c73-42f6-bff3-f3c70f0defa9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "hsi_2013_data shape: (349, 1905, 144)\n",
            "Lidar_2013_data shape: (349, 1905, 1)\n",
            "gt_2013_data.shape: (349, 1905)\n"
          ]
        }
      ],
      "source": [
        "# 2.1 Loads Data\n",
        "# Load hyperpsectral data\n",
        "hsi_2013_data=sio.loadmat(path+'2013_IEEE_GRSS_DF_Contest_CASI_349_1905_144.mat')['ans']\n",
        "print('hsi_2013_data shape:', hsi_2013_data.shape)\n",
        "\n",
        "# Loader Lidar  data\n",
        "import mat73\n",
        "lidar_2013_data = sio.loadmat(path+'2013_IEEE_GRSS_DF_Contest_LiDAR.mat')['LiDAR_data']\n",
        "\n",
        "print('Lidar_2013_data shape:', lidar_2013_data.shape)\n",
        "\n",
        "#Load ground truth labels\n",
        "gt_2013_data=sio.loadmat(path+'GRSS2013.mat')['name']\n",
        "print('gt_2013_data.shape:', gt_2013_data.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ivt40duEwqsP"
      },
      "source": [
        "#1.0  Model Building"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "B6TlJRoRw0z4"
      },
      "outputs": [],
      "source": [
        "# Configuration class\n",
        "class Config:\n",
        "    def __init__(self, in_channels, num_patches, kernel_size, patch_size, emb_size, dim, depth, heads, dim_head, mlp_ratio, num_classes, dropout, pos_emb_size, class_emb_size, stride, output_dim):  # Set default output_dim to 1\n",
        "        self.in_channels = in_channels\n",
        "        self.num_patches = num_patches\n",
        "        self.kernel_size = kernel_size\n",
        "        self.patch_size = patch_size\n",
        "        self.emb_size = emb_size\n",
        "        self.dim = dim\n",
        "        self.depth = depth\n",
        "        self.heads = heads\n",
        "        self.dim_head = dim_head\n",
        "        self.mlp_ratio = mlp_ratio\n",
        "        self.num_classes = num_classes\n",
        "        self.dropout = dropout\n",
        "        self.pos_emb_size = pos_emb_size\n",
        "        self.class_emb_size = class_emb_size\n",
        "        self.stride = stride\n",
        "        self.output_dim = output_dim  # Ensure output_dim is a part of the config"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4a2VjJdHwP1V"
      },
      "source": [
        "### 1.1 Full Architecture Of Forward backward Processing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "ZwRy8jeeXvZ8"
      },
      "outputs": [],
      "source": [
        "# Version 2.0 This involves reversing the input tensor for the backward path before applying the backward_conv1d operation\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "class HSIVimBlock(nn.Module):\n",
        "    def __init__(self, spatial_dim, num_bands, hidden_dim, output_dim, delta_param_init):\n",
        "        super(HSIVimBlock, self).__init__()\n",
        "        # Initialization with self.hidden_dim\n",
        "        self.spatial_dim = spatial_dim\n",
        "        self.num_bands = num_bands\n",
        "        self.hidden_dim = hidden_dim\n",
        "\n",
        "        # LayerNorm is now expecting a flattened feature vector of Bands*H*W elements\n",
        "        self.norm = nn.LayerNorm(num_bands * spatial_dim * spatial_dim)\n",
        "\n",
        "        # Adjust linear layers according to the new input dimension\n",
        "        self.linear_x = nn.Linear(num_bands * spatial_dim * spatial_dim, hidden_dim)\n",
        "        self.linear_z = nn.Linear(num_bands * spatial_dim * spatial_dim, hidden_dim)\n",
        "\n",
        "        self.forward_conv1d = nn.Conv1d(in_channels=hidden_dim, out_channels=hidden_dim, kernel_size=3, padding=1)\n",
        "        self.backward_conv1d = nn.Conv1d(in_channels=hidden_dim, out_channels=hidden_dim, kernel_size=3, padding=1)\n",
        "\n",
        "        self.A = nn.Parameter(torch.randn(hidden_dim, hidden_dim))\n",
        "        self.B = nn.Parameter(torch.randn(hidden_dim, hidden_dim))\n",
        "        #self.C = nn.Parameter(torch.randn(output_dim, hidden_dim))\n",
        "        self.delta_param = nn.Parameter(torch.full((hidden_dim,), delta_param_init))\n",
        "\n",
        "        self.linear_forward = nn.Linear(hidden_dim, output_dim)\n",
        "        self.linear_backward = nn.Linear(hidden_dim, output_dim)\n",
        "\n",
        "    def forward(self, x):\n",
        "        Batch, H, W, Bands = x.shape  # Correct shape extraction assuming [Batch, Height, Width, Bands]\n",
        "\n",
        "        # Correctly reshape for LayerNorm to flatten all spatial and spectral information\n",
        "        x = x.reshape(Batch, -1)  # New shape: [Batch, Bands*H*W]\n",
        "\n",
        "        # Normalize across the flattened spatial-spectral data\n",
        "        x = self.norm(x)\n",
        "\n",
        "        # Projection to hidden dimensions\n",
        "        x_proj = self.linear_x(x)\n",
        "        z_proj = self.linear_z(x)\n",
        "\n",
        "        # Ensure correct reshaping for Conv1d compatibility\n",
        "        x_proj = x_proj.view(Batch, self.hidden_dim, -1)\n",
        "        z_proj = z_proj.view(Batch, self.hidden_dim, -1)\n",
        "\n",
        "        # Reverse z_proj for the backward path\n",
        "        z_proj_reversed = torch.flip(z_proj, dims=[-1])\n",
        "\n",
        "        # Bidirectional Conv1d processing using reversed input for the backward path\n",
        "        x_forward = F.silu(self.forward_conv1d(x_proj))\n",
        "        x_backward = F.silu(self.backward_conv1d(z_proj_reversed))\n",
        "\n",
        "        # Apply delta parameter correctly\n",
        "        delta_expanded = self.delta_param.unsqueeze(0).unsqueeze(2)  # Correct shape for broadcasting\n",
        "\n",
        "        # SSM processing with delta applied, using the original and reversed inputs for forward and backward paths respectively\n",
        "        forward_ssm_output = torch.tanh(self.forward_conv1d(x_proj) + self.A * delta_expanded)\n",
        "        backward_ssm_output = torch.tanh(self.backward_conv1d(z_proj_reversed) + self.B * delta_expanded)\n",
        "\n",
        "        # Combine forward and backward outputs into a single representation\n",
        "        forward_reduced = forward_ssm_output.mean(dim=2)\n",
        "        backward_reduced = backward_ssm_output.mean(dim=2)\n",
        "\n",
        "        # Combine the reduced forward and backward paths\n",
        "        y_forward = self.linear_forward(forward_reduced)\n",
        "        y_backward = self.linear_backward(backward_reduced)\n",
        "\n",
        "        # Element-wise sum of forward and backward outputs\n",
        "        y_combined = y_forward + y_backward\n",
        "\n",
        "        # Return the combined output\n",
        "        return y_combined\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 1.2 S[atialFeature processing"
      ],
      "metadata": {
        "id": "zZnKBx3C9lG0"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "02uU_S6M3EQf"
      },
      "outputs": [],
      "source": [
        "# New version\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "class SpatialFeatureProcessing(nn.Module):\n",
        "    def __init__(self, input_channels):\n",
        "        super(SpatialFeatureProcessing, self).__init__()\n",
        "        self.conv_layers = nn.Sequential(\n",
        "            # First convolutional layer with dilation rate of 1 (standard convolution)\n",
        "            nn.Conv2d(in_channels=input_channels, out_channels=256, kernel_size=(3, 3), padding=1, dilation=1),\n",
        "            nn.ReLU(),\n",
        "            nn.BatchNorm2d(256),\n",
        "            # Second convolutional layer with a higher dilation rate to increase the receptive field\n",
        "            nn.Conv2d(in_channels=256, out_channels=512, kernel_size=(3, 3), padding=2, dilation=2),  # Note the increased padding to maintain the spatial dimensions\n",
        "            nn.ReLU(),\n",
        "            nn.BatchNorm2d(512)\n",
        "        )\n",
        "        self.global_avg_pool = nn.AdaptiveAvgPool2d((1, 1))  # Adding global average pooling\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.conv_layers(x)\n",
        "        x = self.global_avg_pool(x)  # Apply global average pooling\n",
        "        x = torch.flatten(x, start_dim=1)  # Flatten all dimensions except batch\n",
        "        return x\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 1.3 Classifier"
      ],
      "metadata": {
        "id": "lIN7Jlj59qI0"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "DCVCyDOHf0ol"
      },
      "outputs": [],
      "source": [
        "class Classifier(nn.Module):\n",
        "    def __init__(self, in_features, num_classes):\n",
        "        super(Classifier, self).__init__()\n",
        "        self.fc_layers = nn.Sequential(\n",
        "            nn.Linear(in_features=in_features, out_features=1024),\n",
        "            nn.ReLU(),\n",
        "            nn.Dropout(0.5),\n",
        "            nn.Linear(in_features=1024, out_features=num_classes),\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.fc_layers(x)\n",
        "        # Remove softmax here if you're using a loss function that includes it, such as nn.CrossEntropyLoss\n",
        "        return x\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T-iDoAegAj8D"
      },
      "source": [
        "###1.4 Integrated into Main Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "mbvWBj6CjoQm"
      },
      "outputs": [],
      "source": [
        "class HSIClassificationMambaModel(nn.Module):\n",
        "    def __init__(self, spatial_dim, num_bands, hidden_dim, output_dim, delta_param_init, num_classes):\n",
        "        super(HSIClassificationMambaModel, self).__init__()\n",
        "        self.vim_block = HSIVimBlock(spatial_dim, num_bands, hidden_dim, output_dim, delta_param_init)\n",
        "        self.output_dim = output_dim  # Save output_dim as an attribute of the class\n",
        "\n",
        "        # Initialize SpatialFeatureProcessing and Classifier here\n",
        "        # Adjusted to pass 'output_dim' as 'input_channels' to SpatialFeatureProcessing\n",
        "        self.spatial_processing = SpatialFeatureProcessing(input_channels=output_dim)\n",
        "        # Assuming the output of SpatialFeatureProcessing matches the in_features expected by Classifier\n",
        "        self.classifier = Classifier(in_features=512, num_classes=num_classes)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.vim_block(x)\n",
        "        # This is a placeholder. Actual reshaping depends on the output of HSIVimBlock and the input expectation of SpatialFeatureProcessing\n",
        "        x = x.view(-1, self.output_dim, 1, 1)  # Reshape to include spatial dimensions if needed\n",
        "        x = self.spatial_processing(x)\n",
        "\n",
        "        # Flatten the output from spatial processing if it's not already flat\n",
        "        x = torch.flatten(x, start_dim=1)\n",
        "\n",
        "        x = self.classifier(x)\n",
        "        return x\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4YPTp6kBEwgr"
      },
      "source": [
        "# Instance the Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pjaOisQ-e9vP",
        "outputId": "c1cf28c6-6b6e-490c-b34e-2bc141100fe6"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "HSIClassificationMambaModel(\n",
            "  (vim_block): HSIVimBlock(\n",
            "    (norm): LayerNorm((7056,), eps=1e-05, elementwise_affine=True)\n",
            "    (linear_x): Linear(in_features=7056, out_features=256, bias=True)\n",
            "    (linear_z): Linear(in_features=7056, out_features=256, bias=True)\n",
            "    (forward_conv1d): Conv1d(256, 256, kernel_size=(3,), stride=(1,), padding=(1,))\n",
            "    (backward_conv1d): Conv1d(256, 256, kernel_size=(3,), stride=(1,), padding=(1,))\n",
            "    (linear_forward): Linear(in_features=256, out_features=128, bias=True)\n",
            "    (linear_backward): Linear(in_features=256, out_features=128, bias=True)\n",
            "  )\n",
            "  (spatial_processing): SpatialFeatureProcessing(\n",
            "    (conv_layers): Sequential(\n",
            "      (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "      (1): ReLU()\n",
            "      (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "      (3): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2), dilation=(2, 2))\n",
            "      (4): ReLU()\n",
            "      (5): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    )\n",
            "    (global_avg_pool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
            "  )\n",
            "  (classifier): Classifier(\n",
            "    (fc_layers): Sequential(\n",
            "      (0): Linear(in_features=512, out_features=1024, bias=True)\n",
            "      (1): ReLU()\n",
            "      (2): Dropout(p=0.5, inplace=False)\n",
            "      (3): Linear(in_features=1024, out_features=15, bias=True)\n",
            "    )\n",
            "  )\n",
            ")\n"
          ]
        }
      ],
      "source": [
        "\n",
        "# Instantiate the model\n",
        "model = HSIClassificationMambaModel(\n",
        "    spatial_dim=7,\n",
        "    num_bands=144,\n",
        "    hidden_dim=256,\n",
        "    output_dim=128,  # Make sure this matches the actual output of your HSIVimBlock\n",
        "    delta_param_init=0.01,\n",
        "    num_classes=15\n",
        ")\n",
        "\n",
        "# Print the model architecture\n",
        "print(model)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_7hfU1LYg_PW"
      },
      "source": [
        "### Optional 1   Data Preparation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ceR1uP-X1pvx",
        "outputId": "921f6672-ab7d-4499-ceb5-552c2362d990"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{1: {'class_name': 'Healthy grass', 'training_sample': 198, 'test_sample': 1053, 'total_samples': 1251}, 2: {'class_name': 'Stressed grass', 'training_sample': 190, 'test_sample': 1064, 'total_samples': 1254}, 3: {'class_name': 'Synthetic grass', 'training_sample': 192, 'test_sample': 505, 'total_samples': 697}, 4: {'class_name': 'Trees', 'training_sample': 188, 'test_sample': 1058, 'total_samples': 1244}, 5: {'class_name': 'Soil', 'training_sample': 186, 'test_sample': 1056, 'total_samples': 1242}, 6: {'class_name': 'Water', 'training_sample': 182, 'test_sample': 141, 'total_samples': 325}, 7: {'class_name': 'Residential', 'training_sample': 196, 'test_sample': 1072, 'total_samples': 1268}, 8: {'class_name': 'Commercial', 'training_sample': 191, 'test_sample': 1053, 'total_samples': 1244}, 9: {'class_name': 'Road', 'training_sample': 193, 'test_sample': 1059, 'total_samples': 1252}, 10: {'class_name': 'Highway', 'training_sample': 191, 'test_sample': 1036, 'total_samples': 1227}, 11: {'class_name': 'Railway', 'training_sample': 181, 'test_sample': 1054, 'total_samples': 1235}, 12: {'class_name': 'Parking lot 1', 'training_sample': 192, 'test_sample': 1041, 'total_samples': 1233}, 13: {'class_name': 'Parking lot 2', 'training_sample': 184, 'test_sample': 285, 'total_samples': 469}, 14: {'class_name': 'Tennis court', 'training_sample': 181, 'test_sample': 247, 'total_samples': 428}, 15: {'class_name': 'Running track', 'training_sample': 187, 'test_sample': 473, 'total_samples': 660}}\n"
          ]
        }
      ],
      "source": [
        "# 2.1 Define the class information\n",
        "class_info = [(1, \"Healthy grass\", 'training_sample', 198, 'test_sample', 1053,  'total', 1251),\n",
        "    (2, \"Stressed grass\",'training_sample', 190, 'test_sample', 1064,  'total', 1254),\n",
        "    (3, \"Synthetic grass\", 'training_sample', 192, 'test_sample', 505,  'total', 697),\n",
        "    (4, \"Trees\", 'training_sample', 188, 'test_sample', 1058,  'total', 1244),\n",
        "    (5, \"Soil\",'training_sample', 186, 'test_sample', 1056,  'total', 1242),\n",
        "    (6, \"Water\", 'training_sample', 182, 'test_sample', 141,  'total', 325),\n",
        "    (7, \"Residential\", 'training_sample', 196, 'test_sample', 1072,  'total', 1268),\n",
        "    (8, \"Commercial\", 'training_sample', 191, 'test_sample', 1053,  'total', 1244),\n",
        "    (9, \"Road\", 'training_sample', 193, 'test_sample', 1059,  'total', 1252),\n",
        "    (10, \"Highway\", 'training_sample', 191, 'test_sample', 1036,  'total', 1227),\n",
        "    (11, \"Railway\", 'training_sample', 181, 'test_sample', 1054,  'total', 1235),\n",
        "    (12, \"Parking lot 1\", 'training_sample', 192, 'test_sample', 1041,  'total', 1233),\n",
        "    (13, \"Parking lot 2\", 'training_sample', 184, 'test_sample',285,  'total', 469),\n",
        "    (14, \"Tennis court\",'training_sample', 181, 'test_sample', 247,  'total', 428),\n",
        "    (15, \"Running track\", 'training_sample', 187, 'test_sample', 473,  'total', 660)]\n",
        "\n",
        "# Create a dictionary to store class number, class name, and class samples\n",
        "class_dict = {class_number: {\"class_name\": class_name,\n",
        "                             'training_sample': training_sample,\n",
        "                             'test_sample': test_sample,\n",
        "                             \"total_samples\": total}\n",
        "              for class_number, class_name, _, training_sample, _, test_sample, _, total in class_info}\n",
        "\n",
        "print(class_dict)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TUkT_fg8w8uH",
        "outputId": "664596c6-b537-45a5-c49b-62d85f8cbea5"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "hsi_samples shape: (15029, 7, 7, 144)\n",
            "lidar_samples shape: (15029, 7, 7, 1)\n",
            "labels shape: (15029,)\n"
          ]
        }
      ],
      "source": [
        "# 2.2 Samples Extraction\n",
        "\n",
        "# Define patch size and stride\n",
        "patch_size = 7\n",
        "stride = 1\n",
        "\n",
        "# Create an empty list to store patches and labels\n",
        "hsi_samples = []\n",
        "lidar_samples = []\n",
        "labels = []\n",
        "\n",
        "# Initialize a dictionary to store class count\n",
        "class_count = {i: 0 for i in class_dict.keys()}\n",
        "\n",
        "# Function to check if all classes have the required number of samples\n",
        "def all_classes_completed(class_count, class_dict):\n",
        "    return all(class_count[class_num] == class_dict[class_num][\"total_samples\"] for class_num in class_dict.keys())\n",
        "\n",
        "while not all_classes_completed(class_count, class_dict):\n",
        "    # Loop through the ground truth data\n",
        "    for label in class_dict.keys():\n",
        "        # Get the coordinates of the ground truth pixels\n",
        "        #coords = np.argwhere((gt_2013_data == label) & (mask > 0))\n",
        "        coords = np.argwhere(gt_2013_data == label)\n",
        "\n",
        "        # Shuffle the coordinates to randomize the patch extraction\n",
        "        np.random.shuffle(coords)\n",
        "\n",
        "        for coord in coords:\n",
        "            i, j = coord\n",
        "            # Calculate the patch indices\n",
        "            i_start, i_end = i - patch_size // 2, i + patch_size // 2 + 1\n",
        "            j_start, j_end = j - patch_size // 2, j + patch_size // 2 + 1\n",
        "\n",
        "            # Check if the indices are within the bounds of the HSI data\n",
        "            if i_start >= 0 and i_end <= hsi_2013_data.shape[0] and j_start >= 0 and j_end <= hsi_2013_data.shape[1]:\n",
        "                # Extract the patch\n",
        "                hsi_patch = hsi_2013_data[i_start:i_end, j_start:j_end, :]\n",
        "\n",
        "                # Extract the LiDAR patch\n",
        "                lidar_patch = lidar_2013_data[i_start:i_end, j_start:j_end, :]\n",
        "\n",
        "                # If the class count is less than the required samples\n",
        "                if class_count[label] < class_dict[label][\"total_samples\"]:\n",
        "                    # Append the patch and its label to the list\n",
        "                    hsi_samples.append(hsi_patch)\n",
        "                    lidar_samples.append(lidar_patch)\n",
        "                    labels.append(label)\n",
        "                    class_count[label] += 1\n",
        "\n",
        "                    # If all classes have the required number of samples, exit the loop\n",
        "                    if all_classes_completed(class_count, class_dict):\n",
        "                        break\n",
        "\n",
        "# Convert the list of patches and labels into arrays\n",
        "hsi_samples = np.array(hsi_samples)\n",
        "lidar_samples = np.array(lidar_samples)\n",
        "labels = np.array(labels)\n",
        "print('hsi_samples shape:', hsi_samples.shape)\n",
        "print('lidar_samples shape:', lidar_samples.shape)\n",
        "print('labels shape:', labels.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YR4vwscb1NMY",
        "outputId": "e4e27079-425e-49de-8cbc-402112b8bdf6"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "hsi_training_samples shape: (2832, 7, 7, 144)\n",
            "lidar_training_samples shape: (2832, 7, 7, 1)\n",
            "training_labels shape: (2832,)\n",
            "hsi_test_samples shape: (12197, 7, 7, 144)\n",
            "lidar_test_samples shape: (12197, 7, 7, 1)\n",
            "test_labels shape: (12197,)\n"
          ]
        }
      ],
      "source": [
        "# Create training_samples_dict based on class_dict\n",
        "training_samples_dict = {class_num: class_info[\"training_sample\"] for class_num, class_info in class_dict.items()}\n",
        "\n",
        "# Assuming `hsi_samples`, `lidar_samples`, and `labels` have been previously defined\n",
        "# Convert the list of patches and labels into arrays if they aren't already\n",
        "hsi_samples = np.array(hsi_samples)\n",
        "lidar_samples = np.array(lidar_samples)\n",
        "labels = np.array(labels)\n",
        "\n",
        "# Create lists to store training and test samples and labels\n",
        "hsi_training_samples, lidar_training_samples, training_labels = [], [], []\n",
        "hsi_test_samples, lidar_test_samples, test_labels = [], [], []\n",
        "\n",
        "# Split samples into training and test sets based on the desired number of training samples\n",
        "for label, train_samples in training_samples_dict.items():\n",
        "    # Get indices of the current class\n",
        "    class_indices = np.where(labels == label)[0]\n",
        "\n",
        "    # Randomly shuffle the indices\n",
        "    np.random.shuffle(class_indices)\n",
        "\n",
        "    # Split the indices into training and test set indices\n",
        "    train_indices = class_indices[:train_samples]\n",
        "    test_indices = class_indices[train_samples:]\n",
        "\n",
        "    # Add training samples and labels for the current class\n",
        "    hsi_training_samples.extend(hsi_samples[train_indices])\n",
        "    lidar_training_samples.extend(lidar_samples[train_indices])\n",
        "    training_labels.extend(labels[train_indices])\n",
        "\n",
        "    # Add test samples and labels for the current class\n",
        "    hsi_test_samples.extend(hsi_samples[test_indices])\n",
        "    lidar_test_samples.extend(lidar_samples[test_indices])\n",
        "    test_labels.extend(labels[test_indices])\n",
        "\n",
        "# Convert lists back to numpy arrays\n",
        "hsi_training_samples = np.array(hsi_training_samples)\n",
        "lidar_training_samples = np.array(lidar_training_samples)\n",
        "training_labels = np.array(training_labels)\n",
        "\n",
        "hsi_test_samples = np.array(hsi_test_samples)\n",
        "lidar_test_samples = np.array(lidar_test_samples)\n",
        "test_labels = np.array(test_labels)\n",
        "\n",
        "# Print shapes to verify\n",
        "print('hsi_training_samples shape:', hsi_training_samples.shape)\n",
        "print('lidar_training_samples shape:', lidar_training_samples.shape)\n",
        "print('training_labels shape:', training_labels.shape)\n",
        "\n",
        "print('hsi_test_samples shape:', hsi_test_samples.shape)\n",
        "print('lidar_test_samples shape:', lidar_test_samples.shape)\n",
        "print('test_labels shape:', test_labels.shape)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lPqInDBO1LU6",
        "outputId": "07b63714-e9b5-499d-be5f-0734c83b0c6f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Augmented HSI training samples shape: (16992, 7, 7, 144)\n",
            "Augmented LiDAR training samples shape: (16992, 7, 7, 1)\n",
            "Augmented training labels shape: (16992,)\n"
          ]
        }
      ],
      "source": [
        "import numpy as np\n",
        "from scipy.ndimage import rotate\n",
        "\n",
        "def augment_training_data(hsi_training_data, lidar_training_data, training_labels, rotations=[45, 90, 135], flip_up_down=True, flip_left_right=True):\n",
        "    augmented_hsi = []\n",
        "    augmented_lidar = []\n",
        "    augmented_labels = []\n",
        "\n",
        "    for hsi, lidar, label in zip(hsi_training_data, lidar_training_data, training_labels):\n",
        "        # Original data\n",
        "        augmented_hsi.append(hsi)\n",
        "        augmented_lidar.append(lidar)\n",
        "        augmented_labels.append(label)\n",
        "\n",
        "        # Rotations\n",
        "        for angle in rotations:\n",
        "            hsi_rotated = rotate(hsi, angle, axes=(0, 1), reshape=False, mode='nearest')\n",
        "            lidar_rotated = rotate(lidar, angle, axes=(0, 1), reshape=False, mode='nearest')\n",
        "\n",
        "            augmented_hsi.append(hsi_rotated)\n",
        "            augmented_lidar.append(lidar_rotated)\n",
        "            augmented_labels.append(label)\n",
        "\n",
        "        # Flip up-down\n",
        "        if flip_up_down:\n",
        "            hsi_flipped_ud = np.flipud(hsi)\n",
        "            lidar_flipped_ud = np.flipud(lidar)\n",
        "\n",
        "            augmented_hsi.append(hsi_flipped_ud)\n",
        "            augmented_lidar.append(lidar_flipped_ud)\n",
        "            augmented_labels.append(label)\n",
        "\n",
        "        # Flip left-right\n",
        "        if flip_left_right:\n",
        "            hsi_flipped_lr = np.fliplr(hsi)\n",
        "            lidar_flipped_lr = np.fliplr(lidar)\n",
        "\n",
        "            augmented_hsi.append(hsi_flipped_lr)\n",
        "            augmented_lidar.append(lidar_flipped_lr)\n",
        "            augmented_labels.append(label)\n",
        "\n",
        "    return np.array(augmented_hsi), np.array(augmented_lidar), np.array(augmented_labels)\n",
        "\n",
        "# Augmenting the training samples\n",
        "augmented_hsi_training_samples, augmented_lidar_training_samples, augmented_training_labels = augment_training_data(hsi_training_samples, lidar_training_samples, training_labels)\n",
        "\n",
        "# Print shapes to verify the augmented training data\n",
        "print('Augmented HSI training samples shape:', augmented_hsi_training_samples.shape)\n",
        "print('Augmented LiDAR training samples shape:', augmented_lidar_training_samples.shape)\n",
        "print('Augmented training labels shape:', augmented_training_labels.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6aH8AYC73RUK",
        "outputId": "fd0473ea-996a-40db-8b21-bfc474e4adca"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "X_train shape: (15292, 7, 7, 144)\n",
            "X_train_val shape: (1700, 7, 7, 144)\n",
            "y_train shape: (15292,)\n",
            "X_test shape: (12197, 7, 7, 144)\n",
            "y_test shape: (12197,)\n"
          ]
        }
      ],
      "source": [
        "import numpy as np\n",
        "import torch\n",
        "from torch.utils.data import TensorDataset, DataLoader\n",
        "from sklearn.model_selection import train_test_split\n",
        "\n",
        "# Split the augmented training data into training, validationsets\n",
        "X_train, X_val, y_train, y_val = train_test_split(\n",
        "    augmented_hsi_training_samples, augmented_training_labels, test_size=0.1, random_state=42, stratify=augmented_training_labels\n",
        ")\n",
        "X_test=hsi_test_samples\n",
        "y_test=test_labels\n",
        "\n",
        "print('X_train shape:', X_train.shape)\n",
        "print('X_train_val shape:', X_val.shape)\n",
        "print('y_train shape:', y_train.shape)\n",
        "\n",
        "print('X_test shape:', X_test.shape)\n",
        "print('y_test shape:', y_test.shape)\n",
        "\n",
        "\n",
        "# Convert the splitted datasets to tensor datasets\n",
        "train_dataset = TensorDataset(torch.tensor(X_train.astype(np.float32)), torch.tensor(y_train).long())\n",
        "val_dataset = TensorDataset(torch.tensor(X_val.astype(np.float32)), torch.tensor(y_val).long())\n",
        "test_dataset = TensorDataset(torch.tensor(X_test.astype(np.float32)), torch.tensor(y_test).long())\n",
        "\n",
        "# Create DataLoader instances for training, validation, and testing\n",
        "batch_size = 32\n",
        "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)\n",
        "val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)\n",
        "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9YKbJ8f9kYbF"
      },
      "source": [
        "# 5.0 Training Model Memeory and Time calcualtion"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "id": "fhsKBbhgvrUy",
        "outputId": "a77d9ff9-dea5-4ebb-8d1b-112da5b3925f",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Initial Memory Allocated: 45.670912 MB\n"
          ]
        }
      ],
      "source": [
        "# Before the training loop, to record the initial memory usage (GPU)\n",
        "if torch.cuda.is_available():\n",
        "    torch.cuda.reset_peak_memory_stats()  # Reset peak memory stats at the start\n",
        "    initial_memory = torch.cuda.memory_allocated()\n",
        "    print(f\"Initial Memory Allocated: {initial_memory / 1e6} MB\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UtM-XdHTizW-"
      },
      "source": [
        "### 5.1 Training Model for complete forward and backward archtoeture"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "id": "3iQ0dayjXvaN",
        "outputId": "147b0ff9-3273-4a7b-da4e-2efb3cf4f1c9",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch [1/50], Train Loss: 0.0482, Val Loss: 0.0288\n",
            "Validation Loss Decreased(inf--->0.028779) \t Saving The Model\n",
            "Epoch [2/50], Train Loss: 0.0250, Val Loss: 0.0177\n",
            "Validation Loss Decreased(0.028779--->0.017722) \t Saving The Model\n",
            "Epoch [3/50], Train Loss: 0.0175, Val Loss: 0.0139\n",
            "Validation Loss Decreased(0.017722--->0.013879) \t Saving The Model\n",
            "Epoch [4/50], Train Loss: 0.0137, Val Loss: 0.0113\n",
            "Validation Loss Decreased(0.013879--->0.011335) \t Saving The Model\n",
            "Epoch [5/50], Train Loss: 0.0115, Val Loss: 0.0088\n",
            "Validation Loss Decreased(0.011335--->0.008838) \t Saving The Model\n",
            "Epoch [6/50], Train Loss: 0.0097, Val Loss: 0.0074\n",
            "Validation Loss Decreased(0.008838--->0.007444) \t Saving The Model\n",
            "Epoch [7/50], Train Loss: 0.0082, Val Loss: 0.0099\n",
            "Epoch [8/50], Train Loss: 0.0070, Val Loss: 0.0057\n",
            "Validation Loss Decreased(0.007444--->0.005673) \t Saving The Model\n",
            "Epoch [9/50], Train Loss: 0.0063, Val Loss: 0.0060\n",
            "Epoch [10/50], Train Loss: 0.0056, Val Loss: 0.0044\n",
            "Validation Loss Decreased(0.005673--->0.004358) \t Saving The Model\n",
            "Epoch [11/50], Train Loss: 0.0052, Val Loss: 0.0037\n",
            "Validation Loss Decreased(0.004358--->0.003732) \t Saving The Model\n",
            "Epoch [12/50], Train Loss: 0.0048, Val Loss: 0.0048\n",
            "Epoch [13/50], Train Loss: 0.0044, Val Loss: 0.0038\n",
            "Epoch [14/50], Train Loss: 0.0040, Val Loss: 0.0041\n",
            "Epoch [15/50], Train Loss: 0.0037, Val Loss: 0.0045\n",
            "Epoch [16/50], Train Loss: 0.0036, Val Loss: 0.0044\n",
            "Epoch [17/50], Train Loss: 0.0028, Val Loss: 0.0026\n",
            "Validation Loss Decreased(0.003732--->0.002644) \t Saving The Model\n",
            "Epoch [18/50], Train Loss: 0.0031, Val Loss: 0.0056\n",
            "Epoch [19/50], Train Loss: 0.0028, Val Loss: 0.0050\n",
            "Epoch [20/50], Train Loss: 0.0026, Val Loss: 0.0045\n",
            "Epoch [21/50], Train Loss: 0.0025, Val Loss: 0.0035\n",
            "Epoch [22/50], Train Loss: 0.0022, Val Loss: 0.0031\n",
            "Epoch [23/50], Train Loss: 0.0024, Val Loss: 0.0021\n",
            "Validation Loss Decreased(0.002644--->0.002079) \t Saving The Model\n",
            "Epoch [24/50], Train Loss: 0.0022, Val Loss: 0.0027\n",
            "Epoch [25/50], Train Loss: 0.0019, Val Loss: 0.0021\n",
            "Epoch [26/50], Train Loss: 0.0022, Val Loss: 0.0025\n",
            "Epoch [27/50], Train Loss: 0.0018, Val Loss: 0.0014\n",
            "Validation Loss Decreased(0.002079--->0.001368) \t Saving The Model\n",
            "Epoch [28/50], Train Loss: 0.0019, Val Loss: 0.0021\n",
            "Epoch [29/50], Train Loss: 0.0021, Val Loss: 0.0020\n",
            "Epoch [30/50], Train Loss: 0.0018, Val Loss: 0.0072\n",
            "Epoch [31/50], Train Loss: 0.0017, Val Loss: 0.0020\n",
            "Epoch [32/50], Train Loss: 0.0015, Val Loss: 0.0021\n",
            "Epoch [33/50], Train Loss: 0.0014, Val Loss: 0.0030\n",
            "Epoch [34/50], Train Loss: 0.0014, Val Loss: 0.0022\n",
            "Epoch [35/50], Train Loss: 0.0014, Val Loss: 0.0019\n",
            "Epoch [36/50], Train Loss: 0.0013, Val Loss: 0.0012\n",
            "Validation Loss Decreased(0.001368--->0.001215) \t Saving The Model\n",
            "Epoch [37/50], Train Loss: 0.0012, Val Loss: 0.0023\n",
            "Epoch [38/50], Train Loss: 0.0014, Val Loss: 0.0017\n",
            "Epoch [39/50], Train Loss: 0.0011, Val Loss: 0.0018\n",
            "Epoch [40/50], Train Loss: 0.0013, Val Loss: 0.0012\n",
            "Epoch [41/50], Train Loss: 0.0011, Val Loss: 0.0036\n",
            "Epoch [42/50], Train Loss: 0.0012, Val Loss: 0.0016\n",
            "Epoch [43/50], Train Loss: 0.0010, Val Loss: 0.0027\n",
            "Epoch [44/50], Train Loss: 0.0010, Val Loss: 0.0015\n",
            "Epoch [45/50], Train Loss: 0.0012, Val Loss: 0.0027\n",
            "Epoch [46/50], Train Loss: 0.0010, Val Loss: 0.0010\n",
            "Validation Loss Decreased(0.001215--->0.000989) \t Saving The Model\n",
            "Epoch [47/50], Train Loss: 0.0012, Val Loss: 0.0017\n",
            "Epoch [48/50], Train Loss: 0.0009, Val Loss: 0.0011\n",
            "Epoch [49/50], Train Loss: 0.0009, Val Loss: 0.0015\n",
            "Epoch [50/50], Train Loss: 0.0009, Val Loss: 0.0014\n",
            "Finished training. Total training time: 194.05 seconds\n"
          ]
        }
      ],
      "source": [
        "# Training Model in GPU\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "from torch.utils.data import DataLoader\n",
        "import copy\n",
        "import time  # Step 1: Import the time module\n",
        "\n",
        "model = HSIClassificationMambaModel(\n",
        "    spatial_dim=7, num_bands=144, hidden_dim=256, output_dim=128, delta_param_init=0.01, num_classes=15\n",
        ").cuda()\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "optimizer = optim.Adam(model.parameters(), lr=0.00001)\n",
        "\n",
        "epochs = 50\n",
        "best_val_loss = float('inf')\n",
        "best_model_wts = copy.deepcopy(model.state_dict())\n",
        "patience = 10\n",
        "\n",
        "start_time = time.time()  # Step 2: Record the start time\n",
        "\n",
        "for epoch in range(epochs):\n",
        "    model.train()\n",
        "    running_train_loss = 0.0\n",
        "    for inputs, labels in train_loader:\n",
        "        inputs, labels = inputs.cuda(), labels.cuda() # Move the data into CUDA\n",
        "        optimizer.zero_grad()\n",
        "        labels -= 1\n",
        "        outputs = model(inputs)\n",
        "        loss = criterion(outputs, labels)\n",
        "\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        running_train_loss += loss.item()\n",
        "\n",
        "    epoch_train_loss = running_train_loss / len(train_loader.dataset)\n",
        "\n",
        "    model.eval()\n",
        "    val_running_loss = 0.0\n",
        "    with torch.no_grad():\n",
        "        for inputs, labels in val_loader:\n",
        "            inputs, labels = inputs.cuda(), labels.cuda() # Move the data into CUDA\n",
        "            outputs = model(inputs)\n",
        "            labels -= 1\n",
        "            loss = criterion(outputs, labels)\n",
        "            val_running_loss += loss.item()\n",
        "\n",
        "        epoch_val_loss = val_running_loss / len(val_loader.dataset)\n",
        "\n",
        "    print(f'Epoch [{epoch+1}/{epochs}], Train Loss: {epoch_train_loss:.4f}, Val Loss: {epoch_val_loss:.4f}')\n",
        "\n",
        "    if epoch_val_loss < best_val_loss:\n",
        "        print(f'Validation Loss Decreased({best_val_loss:.6f}--->{epoch_val_loss:.6f}) \\t Saving The Model')\n",
        "        best_val_loss = epoch_val_loss\n",
        "        best_model_wts = copy.deepcopy(model.state_dict())\n",
        "        no_improve_epochs = 0\n",
        "    else:\n",
        "        no_improve_epochs += 1\n",
        "\n",
        "    if no_improve_epochs > patience:\n",
        "        print('Early stopping!')\n",
        "        model.load_state_dict(best_model_wts)\n",
        "        break\n",
        "\n",
        "end_time = time.time()  # Step 3: Record the end time\n",
        "total_time = end_time - start_time  # Step 4: Calculate total training time\n",
        "\n",
        "print(f'Finished training. Total training time: {total_time:.2f} seconds')  # Print the total training time"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "id": "qlpOQlNKXvaN",
        "outputId": "76d7186d-ab68-4737-87b0-404bbe32d3ce",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Initial Memory Allocated: 168.244736 MB\n"
          ]
        }
      ],
      "source": [
        "# Before the training loop, to record the initial memory usage (GPU)\n",
        "if torch.cuda.is_available():\n",
        "    torch.cuda.reset_peak_memory_stats()  # Reset peak memory stats at the start\n",
        "    initial_memory = torch.cuda.memory_allocated()\n",
        "    print(f\"Initial Memory Allocated: {initial_memory / 1e6} MB\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U-JV3z-rolSh"
      },
      "source": [
        "### Save the modle"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "id": "3pp-QYHqolSh"
      },
      "outputs": [],
      "source": [
        "# Assuming 'model' is your instance of HSIClassificationModel or any other model\n",
        "# and it's been trained\n",
        "torch.save(model.state_dict(), path+'p7_UH2013_model_state_dict.pth')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dpZ1hNquolSi"
      },
      "source": [
        "### Claculte th test time"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "id": "KKUHZHJpXvaO",
        "outputId": "ab704c07-08bd-4502-9237-61cfc1fb8680",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model saved to /content/drive/MyDrive/A02_RemoteSensingData/UHS_2013_DFTC/p7_Uh2013_model_state_dict.pth\n",
            "Overall Accuracy (OA): 0.9656\n",
            "Average Accuracy (AA): 0.9724\n",
            "Kappa Coefficient: 0.9626\n",
            "Test time: 1.43 seconds\n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "import numpy as np\n",
        "from sklearn.metrics import confusion_matrix, accuracy_score, cohen_kappa_score\n",
        "import time  # Import the time module for timing the test phase\n",
        "\n",
        "# Assuming 'model' is your instance of HSIClassificationModel or any other model\n",
        "# and it's been trained\n",
        "\n",
        "# Save the model\n",
        "model_save_path = path+ 'p7_Uh2013_model_state_dict.pth'\n",
        "torch.save(model.state_dict(), model_save_path)\n",
        "print(f'Model saved to {model_save_path}')\n",
        "\n",
        "# Load the model (make sure to initialize the model architecture first)\n",
        "model.load_state_dict(torch.load(model_save_path))\n",
        "model.to(device)\n",
        "\n",
        "# Ensure the model is in evaluation mode\n",
        "model.eval()\n",
        "\n",
        "# Store predictions and actual labels\n",
        "predictions = []\n",
        "actual_labels = []\n",
        "\n",
        "start_time = time.time()  # Start timing\n",
        "\n",
        "with torch.no_grad():\n",
        "    for hsi_patches, labels in test_loader:\n",
        "        # Move data to the appropriate device\n",
        "        hsi_patches = hsi_patches.to(device)\n",
        "        labels -= 1  # Adjust labels if necessary\n",
        "\n",
        "        # Forward pass\n",
        "        outputs = model(hsi_patches)\n",
        "\n",
        "        # Get predictions\n",
        "        _, predicted = torch.max(outputs, 1)\n",
        "        predictions.extend(predicted.cpu().numpy())\n",
        "        actual_labels.extend(labels.cpu().numpy())\n",
        "\n",
        "end_time = time.time()  # End timing\n",
        "test_time = end_time - start_time  # Calculate the test time\n",
        "\n",
        "# Optionally, calculate accuracy or other metrics using predictions and actual_labels\n",
        "\n",
        "# Convert lists to NumPy arrays for easier manipulation\n",
        "predictions_array = np.array(predictions)\n",
        "actual_labels_array = np.array(actual_labels)\n",
        "\n",
        "# Overall Accuracy\n",
        "oa = accuracy_score(actual_labels_array, predictions_array)\n",
        "\n",
        "# Confusion Matrix\n",
        "cm = confusion_matrix(actual_labels_array, predictions_array)\n",
        "# Calculate per-class accuracy from the confusion matrix\n",
        "class_accuracy = cm.diagonal() / cm.sum(axis=1)\n",
        "# Average Accuracy\n",
        "aa = np.mean(class_accuracy)\n",
        "\n",
        "# Kappa Coefficient\n",
        "kappa = cohen_kappa_score(actual_labels_array, predictions_array)\n",
        "\n",
        "print(f'Overall Accuracy (OA): {oa:.4f}')\n",
        "print(f'Average Accuracy (AA): {aa:.4f}')\n",
        "print(f'Kappa Coefficient: {kappa:.4f}')\n",
        "print(f'Test time: {test_time:.2f} seconds')  # Print the test time"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "a8IekWMVolSi",
        "outputId": "dffb26ae-88e5-4792-d267-ad3d4d50008e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Class 1 Accuracy: 0.9962\n",
            "Class 2 Accuracy: 0.9925\n",
            "Class 3 Accuracy: 1.0000\n",
            "Class 4 Accuracy: 0.9792\n",
            "Class 5 Accuracy: 0.9991\n",
            "Class 6 Accuracy: 1.0000\n",
            "Class 7 Accuracy: 0.9757\n",
            "Class 8 Accuracy: 0.9839\n",
            "Class 9 Accuracy: 0.9235\n",
            "Class 10 Accuracy: 0.9701\n",
            "Class 11 Accuracy: 0.8397\n",
            "Class 12 Accuracy: 0.9472\n",
            "Class 13 Accuracy: 0.9789\n",
            "Class 14 Accuracy: 1.0000\n",
            "Class 15 Accuracy: 1.0000\n"
          ]
        }
      ],
      "source": [
        "for i, acc in enumerate(class_accuracy): print(f'Class {i+1} Accuracy: {acc:.4f}')\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "A100",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}