{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":3004,"databundleVersionId":861823,"sourceType":"competition"}],"dockerImageVersionId":31012,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-04-13T15:42:34.164915Z","iopub.execute_input":"2025-04-13T15:42:34.165363Z","iopub.status.idle":"2025-04-13T15:42:36.240226Z","shell.execute_reply.started":"2025-04-13T15:42:34.165331Z","shell.execute_reply":"2025-04-13T15:42:36.239382Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install mamba-ssm","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-14T16:42:23.526703Z","iopub.execute_input":"2025-04-14T16:42:23.526866Z","iopub.status.idle":"2025-04-14T16:45:44.168249Z","shell.execute_reply.started":"2025-04-14T16:42:23.526851Z","shell.execute_reply":"2025-04-14T16:45:44.167092Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nfrom mamba_ssm import Mamba\n\nbatch, length, dim = 2, 64, 16\nx = torch.randn(batch, length, dim).to(\"cuda\")\nmodel = Mamba(\n    # This module uses roughly 3 * expand * d_model^2 parameters\n    d_model=dim, # Model dimension d_model\n    d_state=16,  # SSM state expansion factor\n    d_conv=4,    # Local convolution width\n    expand=2,    # Block expansion factor\n).to(\"cuda\")\ny = model(x)\nassert y.shape == x.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-14T16:45:44.172848Z","iopub.execute_input":"2025-04-14T16:45:44.173113Z","iopub.status.idle":"2025-04-14T16:45:51.836946Z","shell.execute_reply.started":"2025-04-14T16:45:44.173082Z","shell.execute_reply":"2025-04-14T16:45:51.836367Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom mamba_ssm import Mamba\n\n# Mamba Model\nclass MNISTMambaClassifier(nn.Module):\n    def __init__(self, d_model=64, d_state=16, num_classes=10):\n        super().__init__()\n        self.in_proj = nn.Linear(1, d_model)\n        self.mamba = Mamba(d_model=d_model, d_state=d_state, expand=2, d_conv=56)\n        self.fc = nn.Linear(d_model, num_classes)\n    \n    def forward(self, x):\n        x = self.in_proj(x)\n        \n        # Forward propagate Mamba\n        out = self.mamba(x)\n        \n        # Decode the hidden state of the last time step\n        out = self.fc(out[:, -1, :])\n        return out","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-14T17:26:50.222751Z","iopub.execute_input":"2025-04-14T17:26:50.223387Z","iopub.status.idle":"2025-04-14T17:26:50.228762Z","shell.execute_reply.started":"2025-04-14T17:26:50.223365Z","shell.execute_reply":"2025-04-14T17:26:50.227960Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport pandas as pd\nfrom sklearn.model_selection import train_test_split\nimport numpy as np\n\nUSE_MLP = False\n\n# Custom Dataset class\nclass MNISTSequenceDataset(Dataset):\n    def __init__(self, data, labels, mean=0.1307, std=0.3081):\n        # First normalize to [0,1] by dividing by 255\n        normalized_data = data.astype(np.float32) / 255.0\n        \n        # Then apply mean/std normalization\n        normalized_data = (normalized_data - mean) / std\n        \n        self.data = torch.FloatTensor(normalized_data)\n        self.labels = torch.LongTensor(labels)\n    \n    def __len__(self):\n        return len(self.labels)\n    \n    def __getitem__(self, idx):\n        return self.data[idx], self.labels[idx]\n\n# LSTM Model\nclass MNISTLSTMClassifier(nn.Module):\n    def __init__(self, input_size=1, hidden_size=128, num_layers=1, num_classes=10):\n        super().__init__()\n        self.hidden_size = hidden_size\n        self.num_layers = num_layers\n        \n        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)\n        self.fc = nn.Linear(hidden_size, num_classes)\n    \n    def forward(self, x):\n        # Initialize hidden state with zeros\n        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)\n        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)\n        \n        # Forward propagate LSTM\n        out, _ = self.lstm(x, (h0, c0))\n        \n        # Decode the hidden state of the last time step\n        out = self.fc(out[:, -1, :])\n        return out\n\n# MLP Model\nclass MNISTMLPClassifier(nn.Module):\n    def __init__(self, input_size=784, hidden_size=128, num_classes=10):\n        super().__init__()\n        self.input_size = input_size\n        self.num_classes = num_classes\n        self.model = nn.Sequential(\n            nn.Linear(input_size, hidden_size),\n            nn.ReLU(),\n            nn.Dropout(0.2),\n            nn.Linear(hidden_size, hidden_size),\n            nn.ReLU(),\n            nn.Dropout(0.2),\n            nn.Linear(hidden_size, num_classes)\n        )\n    \n    def forward(self, x):\n        x = x.view(-1, self.input_size)\n        return self.model(x)\n\n# Training function\ndef train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs, device):\n    train_losses = []\n    val_losses = []\n    train_accs = []\n    val_accs = []\n    for epoch in range(num_epochs):\n        model.train()\n        train_loss = 0\n        train_correct = 0\n        train_total = 0\n        \n        for images, labels in train_loader:\n            images = images.to(device)\n            labels = labels.to(device)\n            \n            # Reshape images to sequence format (batch_size, sequence_length, input_size)\n            images = images.view(-1, 784, 1)\n            \n            # Forward pass\n            outputs = model(images)\n            loss = criterion(outputs, labels)\n            \n            # Backward and optimize\n            optimizer.zero_grad()\n            loss.backward()\n            optimizer.step()\n            \n            train_loss += loss.item()\n            _, predicted = torch.max(outputs.data, 1)\n            train_total += labels.size(0)\n            train_correct += (predicted == labels).sum().item()\n\n        train_losses.append(train_loss)\n        train_accs.append(train_correct / train_total)\n        \n        # Validation\n        model.eval()\n        val_loss = 0\n        val_correct = 0\n        val_total = 0\n        \n        with torch.no_grad():\n            for images, labels in val_loader:\n                images = images.to(device)\n                labels = labels.to(device)\n                images = images.view(-1, 784, 1)\n                \n                outputs = model(images)\n                loss = criterion(outputs, labels)\n                \n                val_loss += loss.item()\n                _, predicted = torch.max(outputs.data, 1)\n                val_total += labels.size(0)\n                val_correct += (predicted == labels).sum().item()\n\n            val_losses.append(val_loss)\n            val_accs.append(val_correct / val_total)\n        \n        print(f'Epoch [{epoch+1}/{num_epochs}]')\n        print(f'Train Loss: {train_loss/len(train_loader):.4f}, Train Acc: {100*train_correct/train_total:.2f}%')\n        print(f'Val Loss: {val_loss/len(val_loader):.4f}, Val Acc: {100*val_correct/val_total:.2f}%')\n        print('--------------------')\n\n    return train_losses, val_losses, train_accs, val_accs\n\n\n# Set device\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\n# Load and prepare data\ndf = pd.read_csv('/kaggle/input/digit-recognizer/train.csv')\nlabels = df['label'].values\npixels = df.drop('label', axis=1).values\n\n# Split the data\nX_train, X_val, y_train, y_val = train_test_split(pixels, labels, test_size=0.2, random_state=42)\n\n# Create datasets\ntrain_dataset = MNISTSequenceDataset(X_train, y_train)\nval_dataset = MNISTSequenceDataset(X_val, y_val)\n\n# Create data loaders\ntrain_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)\nval_loader = DataLoader(val_dataset, batch_size=256, shuffle=False)\n\n# Initialize model\nif USE_MLP:\n    model = MNISTMLPClassifier().to(device)\nelse:\n    model = MNISTMambaClassifier(d_model=64).to(device)\n\n# Loss and optimizer\ncriterion = nn.CrossEntropyLoss()\noptimizer = torch.optim.RMSprop(model.parameters(), lr=0.001)\n\n# Train the model\nnum_epochs = 50\ntrain_losses, val_losses, train_accs, val_accs = train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs, device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-14T17:27:10.229338Z","iopub.execute_input":"2025-04-14T17:27:10.230087Z","iopub.status.idle":"2025-04-14T17:43:19.149361Z","shell.execute_reply.started":"2025-04-14T17:27:10.230061Z","shell.execute_reply":"2025-04-14T17:43:19.148647Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\n# Create epochs array for x-axis\nepochs = range(1, len(train_losses) + 1)\n\n# Create figure with two subplots\nplt.figure(figsize=(12, 4))\n\n# Plot losses\nplt.subplot(1, 2, 1)\nplt.plot(epochs, train_losses, 'b-', label='Training Loss')\nplt.plot(epochs, val_losses, 'r-', label='Validation Loss')\nplt.title('Training and Validation Loss')\nplt.xlabel('Epochs')\nplt.ylabel('Loss')\nplt.legend()\n\n# Plot accuracies\nplt.subplot(1, 2, 2)\nplt.plot(epochs, train_accs, 'b-', label='Training Accuracy')\nplt.plot(epochs, val_accs, 'r-', label='Validation Accuracy')\nplt.title('Training and Validation Accuracy')\nplt.xlabel('Epochs')\nplt.ylabel('Accuracy')\nplt.legend()\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-14T17:43:19.150410Z","iopub.execute_input":"2025-04-14T17:43:19.150672Z","iopub.status.idle":"2025-04-14T17:43:19.500040Z","shell.execute_reply.started":"2025-04-14T17:43:19.150656Z","shell.execute_reply":"2025-04-14T17:43:19.499205Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = MNISTMambaClassifier(d_model=64)\n\ntotal = 0\nfor name, param in model.named_parameters():\n    print(f\"{name}: {param.numel()}\")\n    total += param.numel()\nprint(total)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-14T18:13:26.151216Z","iopub.execute_input":"2025-04-14T18:13:26.151903Z","iopub.status.idle":"2025-04-14T18:13:26.160572Z","shell.execute_reply.started":"2025-04-14T18:13:26.151877Z","shell.execute_reply":"2025-04-14T18:13:26.159813Z"}},"outputs":[],"execution_count":null}]}