{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[],"dockerImageVersionId":28755,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# MNIST Classification - Optuna, Model Ensemble\n## Detailed guide into training and evaluating a model for MNIST Classification for high accuracy","metadata":{}},{"cell_type":"markdown","source":"This notebook intends to provide a comprehensive guide to achieving a high accuracy in the MNIST dataset. Data loading, augmentation and model creation will happen via pytorch, hyperparameters will be optimized using optuna search, and the top 4 models after repeated runs will create a model ensemble for final predictions","metadata":{}},{"cell_type":"markdown","source":"## Imports","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport matplotlib.pyplot as plt\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\n\nfrom sklearn.model_selection import train_test_split \n\nimport optuna","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T17:35:11.115341Z","iopub.execute_input":"2026-08-25T17:35:11.115656Z","iopub.status.idle":"2026-08-25T17:35:17.369233Z","shell.execute_reply.started":"2026-08-25T17:35:11.115628Z","shell.execute_reply":"2026-08-25T17:35:17.368609Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Load data","metadata":{}},{"cell_type":"code","source":"#loading images into pandas dataframe\ntrain = pd.read_csv(\"/kaggle/input/competitions/digit-recognizer/train.csv\")\ntest = pd.read_csv(\"/kaggle/input/competitions/digit-recognizer/test.csv\")\n\n# Assigning X as training data, and Y as target data\n# Reshape into (no of training examples, 28 x 28 image, 1 channel)\n\nX_train = train.drop(\"label\", axis=1).values.reshape(-1, 1, 28, 28)\ny_train = train[\"label\"].values\nX_test = test.values.reshape(-1, 1, 28, 28)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T17:35:17.370872Z","iopub.execute_input":"2026-08-25T17:35:17.371187Z","iopub.status.idle":"2026-08-25T17:35:21.893262Z","shell.execute_reply.started":"2026-08-25T17:35:17.371166Z","shell.execute_reply":"2026-08-25T17:35:21.892632Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Display Data","metadata":{}},{"cell_type":"code","source":"plt.figure(figsize=(18,6))\n\nfor i in range(36):\n    plt.subplot(3,12,i+1)\n\n    image = X_train[i].squeeze()\n\n    plt.imshow(image,cmap=\"binary\")\n    plt.title(y_train[i])\n    plt.axis(\"off\")\n\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T17:35:21.894078Z","iopub.execute_input":"2026-08-25T17:35:21.894461Z","iopub.status.idle":"2026-08-25T17:35:23.086921Z","shell.execute_reply.started":"2026-08-25T17:35:21.894373Z","shell.execute_reply":"2026-08-25T17:35:23.085974Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Splitting Data","metadata":{}},{"cell_type":"code","source":"X_train, X_dev, y_train, y_dev = train_test_split(X_train, y_train, test_size = 0.1, random_state = 42, stratify = y_train)\n\n# Creates a validation set for hyperparam optimization with split 0.9/0.1, feel free to try different splits \n# Stratify = y_train ensures the labels distribution mantains consistent across y_train and y_dev","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T17:35:23.088116Z","iopub.execute_input":"2026-08-25T17:35:23.088472Z","iopub.status.idle":"2026-08-25T17:35:23.497599Z","shell.execute_reply.started":"2026-08-25T17:35:23.088446Z","shell.execute_reply":"2026-08-25T17:35:23.496716Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fig, axes = plt.subplots(1,2,figsize=(12,4))\n\nfor ax,y,title in zip(axes,[y_train,y_dev],[\"Train\",\"Dev\"]):\n    pd.Series(y).value_counts().sort_index().plot.bar(ax=ax,title=title)\n\nplt.show()\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T17:35:23.498585Z","iopub.execute_input":"2026-08-25T17:35:23.498964Z","iopub.status.idle":"2026-08-25T17:35:23.801261Z","shell.execute_reply.started":"2026-08-25T17:35:23.498905Z","shell.execute_reply":"2026-08-25T17:35:23.800501Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Normalizing data","metadata":{}},{"cell_type":"code","source":"mean = X_train.mean()\nstd = X_train.std()\n\nX_train = (X_train - mean) / std\nX_dev = (X_dev - mean) / std\nX_test = (X_test - mean) / std\n\n#Fit on train set and then apply the same transformation to dev set","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T17:35:23.802079Z","iopub.execute_input":"2026-08-25T17:35:23.802521Z","iopub.status.idle":"2026-08-25T17:35:24.196101Z","shell.execute_reply.started":"2026-08-25T17:35:23.802495Z","shell.execute_reply":"2026-08-25T17:35:24.195384Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Data augmentation","metadata":{}},{"cell_type":"code","source":"from torchvision import transforms\n\ntrain_transform = transforms.Compose([\n    transforms.RandomRotation(10),\n    transforms.RandomAffine(\n        degrees=0,\n        translate=(0.1,0.1)\n    )\n])\n\n#randomly rotate up to 10 degrees, and shift up to 10% of the image size in either direction","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T17:35:24.198330Z","iopub.execute_input":"2026-08-25T17:35:24.198711Z","iopub.status.idle":"2026-08-25T17:35:29.739061Z","shell.execute_reply.started":"2026-08-25T17:35:24.198684Z","shell.execute_reply":"2026-08-25T17:35:29.738181Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Vizualization the transformations\n\nimage = torch.tensor(X_train[0], dtype=torch.float32)\n\nfig, axes = plt.subplots(2,4,figsize=(12,4))\n\nfor ax in axes.flat:\n    ax.imshow(train_transform(image).squeeze(), cmap=\"gray\")\n    ax.axis(\"off\")\n\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T17:35:29.740010Z","iopub.execute_input":"2026-08-25T17:35:29.740413Z","iopub.status.idle":"2026-08-25T17:35:29.940074Z","shell.execute_reply.started":"2026-08-25T17:35:29.740386Z","shell.execute_reply":"2026-08-25T17:35:29.939407Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Dataset Class","metadata":{}},{"cell_type":"code","source":"class MNISTDataset(Dataset):\n    def __init__(self, images, labels=None, transform=None):\n        self.images = torch.tensor(images, dtype=torch.float32)\n        self.labels = labels\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.images)\n\n    def __getitem__(self, index):\n        image = self.images[index]\n\n        if self.transform:\n            image = self.transform(image)\n\n        if self.labels is None:\n            return image\n\n        label = torch.tensor(self.labels[index], dtype=torch.long)\n        return image, label","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T17:35:29.940937Z","iopub.execute_input":"2026-08-25T17:35:29.941330Z","iopub.status.idle":"2026-08-25T17:35:29.948005Z","shell.execute_reply.started":"2026-08-25T17:35:29.941303Z","shell.execute_reply":"2026-08-25T17:35:29.947367Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_dataset = MNISTDataset(X_train, y_train, transform=train_transform)\ndev_dataset = MNISTDataset(X_dev, y_dev)\ntest_dataset = MNISTDataset(X_test)\n\ntrain_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)\ndev_loader = DataLoader(dev_dataset, batch_size=64, shuffle=False)\ntest_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T17:35:29.948908Z","iopub.execute_input":"2026-08-25T17:35:29.949247Z","iopub.status.idle":"2026-08-25T17:35:30.060259Z","shell.execute_reply.started":"2026-08-25T17:35:29.949223Z","shell.execute_reply":"2026-08-25T17:35:30.059620Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"## Preview the shapes \nimages, labels = next(iter(train_loader))\n\nprint(\"Images shape:\", images.shape)\nprint(\"Labels shape:\", labels.shape)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T17:35:30.061167Z","iopub.execute_input":"2026-08-25T17:35:30.061571Z","iopub.status.idle":"2026-08-25T17:35:30.134248Z","shell.execute_reply.started":"2026-08-25T17:35:30.061523Z","shell.execute_reply":"2026-08-25T17:35:30.133301Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Flexible Architechture","metadata":{}},{"cell_type":"markdown","source":"The architecture below allows customization of hyperparameters - filter sizes, no of layers, kernel size, dropout_rate, neurons in fc layer, batchnorm true or false \nYou can easily add more hyperparameters to explore by setting them as variables","metadata":{}},{"cell_type":"code","source":"class MNISTClassifier(nn.Module):\n    def __init__(self, n_filters, n_layers, kernel_size, dropout_rate, fc_size, use_batchnorm):\n        super().__init__()\n\n        blocks = []\n        in_channels = 1\n        padding = (kernel_size - 1) // 2\n\n        for i in range(n_layers):\n            out_channels = n_filters[i]\n\n            layers = [\n                nn.Conv2d(\n                    in_channels,\n                    out_channels,\n                    kernel_size=kernel_size,\n                    padding=padding\n                )\n            ]\n\n            if use_batchnorm:\n                layers.append(nn.BatchNorm2d(out_channels))\n\n            layers.extend([\n                nn.ReLU(),\n                nn.MaxPool2d(2)\n            ])\n\n            blocks.append(nn.Sequential(*layers))\n            in_channels = out_channels\n\n        self.features = nn.Sequential(*blocks)\n\n        spatial_size = 28 // (2 ** n_layers)\n\n        self.classifier = nn.Sequential(\n            nn.Flatten(),\n            nn.Dropout(dropout_rate),\n            nn.Linear(\n                n_filters[-1] * spatial_size * spatial_size,\n                fc_size\n            ),\n            nn.ReLU(),\n            nn.Dropout(dropout_rate),\n            nn.Linear(fc_size, 10)\n        )\n\n    def forward(self, x):\n        x = self.features(x)\n        x = self.classifier(x)\n        return x","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T17:35:30.135337Z","iopub.execute_input":"2026-08-25T17:35:30.135648Z","iopub.status.idle":"2026-08-25T17:35:30.144697Z","shell.execute_reply.started":"2026-08-25T17:35:30.135623Z","shell.execute_reply":"2026-08-25T17:35:30.143975Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Training","metadata":{}},{"cell_type":"markdown","source":"Below is a sample training run with optimized hyperparameters pre put in.\nWe will explore how to find these optimized hyperparameters through optuna later","metadata":{}},{"cell_type":"code","source":"\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nmodel = MNISTClassifier(\n    n_filters=[64,128,256,256],\n    n_layers=4,\n    kernel_size=5,\n    dropout_rate=0.2,\n    fc_size=256,\n    use_batchnorm=True\n).to(device)\n\noptimizer = torch.optim.AdamW(model.parameters(),lr=0.001,weight_decay=1e-4)\n\nscheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,T_max=20)\n\nloss_function = nn.CrossEntropyLoss()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T17:35:30.145567Z","iopub.execute_input":"2026-08-25T17:35:30.145891Z","iopub.status.idle":"2026-08-25T17:35:30.723042Z","shell.execute_reply.started":"2026-08-25T17:35:30.145866Z","shell.execute_reply":"2026-08-25T17:35:30.722344Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Typical pytorch training loop\ndef train_epoch(model,loader,optimizer,loss_function,device):\n    model.train()\n    total_loss = 0\n\n    for images,labels in loader:\n        images,labels = images.to(device),labels.to(device)\n\n        optimizer.zero_grad()\n        outputs = model(images)\n        loss = loss_function(outputs,labels)\n\n        loss.backward()\n        optimizer.step()\n\n        total_loss += loss.item()\n\n    return total_loss / len(loader)\n\nfor epoch in range(5):\n    train_loss = train_epoch(model,train_loader,optimizer,loss_function,device)\n    print(f\"Epoch {epoch+1}/10 | Loss: {train_loss:.4f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T17:41:21.507994Z","iopub.execute_input":"2026-08-25T17:41:21.508386Z","iopub.status.idle":"2026-08-25T17:45:33.185975Z","shell.execute_reply.started":"2026-08-25T17:41:21.508359Z","shell.execute_reply":"2026-08-25T17:45:33.185145Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Evaluate accuracy","metadata":{}},{"cell_type":"code","source":"# Create a function to evaluate accuracy\ndef evaluate(model, loader, device):\n    model.eval()\n    correct = 0\n    total = 0\n\n    with torch.no_grad():\n        for images, labels in loader:\n            images, labels = images.to(device), labels.to(device)\n\n            predictions = model(images).argmax(dim=1)\n\n            total += labels.size(0)\n            correct += (predictions == labels).sum().item()\n\n    return correct / total * 100\n\naccuracy = evaluate(model,dev_loader,device)\nprint(f\"Dev accuracy is {accuracy}%\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T17:49:00.288138Z","iopub.execute_input":"2026-08-25T17:49:00.288793Z","iopub.status.idle":"2026-08-25T17:49:00.815369Z","shell.execute_reply.started":"2026-08-25T17:49:00.288762Z","shell.execute_reply":"2026-08-25T17:49:00.814480Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Optuna Study","metadata":{}},{"cell_type":"markdown","source":"Optuna randomly selects hyperparams within predefined bounds. It searches more densely in areas that have previously led to a higher accuracy. You can experiment by changing the bounds, or rerunning the study when you have found tighter boundaries for hyperparameters","metadata":{}},{"cell_type":"code","source":"# Create the Optuna Study\ndef objective(trial):\n    n_layers = trial.suggest_int(\"n_layers\",2,4)\n    initial_filters = trial.suggest_categorical(\"initial_filters\",[16,32,64])\n    kernel_size = trial.suggest_categorical(\"kernel_size\",[3,5,7])\n    use_batchnorm = trial.suggest_categorical(\"use_batchnorm\",[True,False])\n    fc_size = trial.suggest_categorical(\"fc_size\",[32,64,128,256])\n    n_filters = [min(initial_filters * (2 ** i),256) for i in range(n_layers)]\n\n    print(f\"\\nTrial {trial.number}\")\n    print(f\"layers={n_layers} | filters={n_filters} | kernel={kernel_size} | batchnorm={use_batchnorm} | fc={fc_size}\")\n\n    model = MNISTClassifier(\n        n_layers=n_layers,\n        n_filters=n_filters,\n        kernel_size=kernel_size,\n        dropout_rate=0.2,\n        fc_size=fc_size,\n        use_batchnorm=use_batchnorm\n    ).to(device)\n\n    optimizer = torch.optim.AdamW(model.parameters(),lr=0.001,weight_decay=1e-4)\n    loss_function = nn.CrossEntropyLoss()\n\n    for epoch in range(5):\n        train_epoch(model,train_loader,optimizer,loss_function,device)\n\n    accuracy = evaluate(model,dev_loader,device)\n\n    print(f\"accuracy={accuracy:.2f}%\")\n    return accuracy\n\nstudy = optuna.create_study(direction=\"maximize\")\nstudy.optimize(objective,n_trials=5)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T18:01:44.792100Z","iopub.execute_input":"2026-08-25T18:01:44.792639Z","iopub.status.idle":"2026-08-25T18:18:48.082634Z","shell.execute_reply.started":"2026-08-25T18:01:44.792605Z","shell.execute_reply":"2026-08-25T18:18:48.081843Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# View the study results \nresults = study.trials_dataframe()\nresults.sort_values(\"value\",ascending=False).head(10)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T18:18:57.889860Z","iopub.execute_input":"2026-08-25T18:18:57.890283Z","iopub.status.idle":"2026-08-25T18:18:57.933144Z","shell.execute_reply.started":"2026-08-25T18:18:57.890250Z","shell.execute_reply":"2026-08-25T18:18:57.932444Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Model Ensemble","metadata":{}},{"cell_type":"markdown","source":"Train the best architecture multiple times, since random initialization changes accuracy. Save the top 3 models and average their probability vectors in each test input to reduce random errors.More runs and more epochs (saving on best epochs) will typically lead to a higher accuracy. I personally achieved 99.66% accuracy after numerous runs and an esemble of 5 models\n\nThe hyperparameters used in the model below were found through optuna, you may change it to match your best found hyperparameters ","metadata":{}},{"cell_type":"code","source":"# Run 20 times, get the 3 models with highest accuracy\nimport torch\nimport torch.nn as nn\nimport copy\n\ntop_models = []\n\nfor run in range(5):\n    print(f\"\\nRun {run+1}/5\")\n\n    model = MNISTClassifier(\n        n_layers=4,\n        n_filters=[64,128,256,256],\n        kernel_size=5,\n        dropout_rate=0.2,\n        fc_size=256,\n        use_batchnorm=True\n    ).to(device)\n\n    optimizer = torch.optim.AdamW(model.parameters(),lr=0.001,weight_decay=1e-4)\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,T_max=10)\n    loss_function = nn.CrossEntropyLoss()\n\n    best_accuracy = 0\n    best_state = None\n    best_epoch = 0\n\n    for epoch in range(10):\n        loss = train_epoch(model,train_loader,optimizer,loss_function,device)\n        accuracy = evaluate(model,dev_loader,device)\n\n        print(f\"Epoch {epoch+1}/10 | Loss: {loss:.4f} | Accuracy: {accuracy:.2f}%\")\n\n        if accuracy > best_accuracy:\n            best_accuracy = accuracy\n            best_epoch = epoch + 1\n            best_state = copy.deepcopy(model.state_dict())\n\n        scheduler.step()\n\n    top_models.append({\n        \"model_state_dict\": best_state,\n        \"accuracy\": best_accuracy,\n        \"epoch\": best_epoch,\n        \"run\": run + 1\n    })\n\n    top_models = sorted(top_models,key=lambda x:x[\"accuracy\"],reverse=True)[:3]\n\n    for i,checkpoint in enumerate(top_models):\n        torch.save(checkpoint,f\"bestmodel_{i+1}.pth\")\n\n    print(f\"Best for run {run+1}: {best_accuracy:.2f}% at epoch {best_epoch}\")\n    print(\"Current top 3:\",[round(m[\"accuracy\"],2) for m in top_models])\n\nprint(\"\\nFinal top 3:\")\nfor i,checkpoint in enumerate(top_models):\n    print(f\"bestmodel_{i+1}.pth | {checkpoint['accuracy']:.2f}% | Run {checkpoint['run']} | Epoch {checkpoint['epoch']}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T18:25:58.655107Z","iopub.execute_input":"2026-08-25T18:25:58.655719Z","iopub.status.idle":"2026-08-25T19:07:36.632914Z","shell.execute_reply.started":"2026-08-25T18:25:58.655689Z","shell.execute_reply":"2026-08-25T19:07:36.631910Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load Each Model\ndef load_model(path):\n    checkpoint = torch.load(path,map_location=device)\n\n    model = MNISTClassifier(\n        n_layers=4,\n        n_filters=[64,128,256,256],\n        kernel_size=5,\n        dropout_rate=0.2,\n        fc_size=256,\n        use_batchnorm=True\n    ).to(device)\n\n    model.load_state_dict(checkpoint[\"model_state_dict\"])\n    model.eval()\n\n    return model\nmodels = [\n    load_model(\"/kaggle/working/bestmodel_1.pth\"),\n    load_model(\"/kaggle/working/bestmodel_2.pth\"),\n    load_model(\"/kaggle/working/bestmodel_3.pth\")\n]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T19:09:32.853691Z","iopub.execute_input":"2026-08-25T19:09:32.854290Z","iopub.status.idle":"2026-08-25T19:09:32.969105Z","shell.execute_reply.started":"2026-08-25T19:09:32.854258Z","shell.execute_reply":"2026-08-25T19:09:32.968434Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Run each model on the dev images, and average the results to predict the number\npredictions = []\nwrong_images = []\nwrong_preds = []\nwrong_labels = []\ncorrect = total = 0\n\nwith torch.no_grad():\n    for images,labels in dev_loader:\n        images,labels = images.to(device),labels.to(device)\n\n        probs = sum(torch.softmax(model(images),dim=1) for model in models) / len(models)\n        predicted = probs.argmax(dim=1)\n\n        predictions.extend(predicted.cpu().numpy())\n        total += labels.size(0)\n        correct += (predicted == labels).sum().item()\n\n        wrong = predicted != labels\n        wrong_images.extend(images[wrong].cpu())\n        wrong_preds.extend(predicted[wrong].cpu())\n        wrong_labels.extend(labels[wrong].cpu())\n\naccuracy = correct / total * 100\nprint(f\"Accuracy: {accuracy:.2f}%\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T19:09:35.588434Z","iopub.execute_input":"2026-08-25T19:09:35.588851Z","iopub.status.idle":"2026-08-25T19:09:36.924069Z","shell.execute_reply.started":"2026-08-25T19:09:35.588824Z","shell.execute_reply":"2026-08-25T19:09:36.923440Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## View wrong Predicted Images","metadata":{}},{"cell_type":"code","source":"fig,axes = plt.subplots(2,5,figsize=(8,3))\n\nfor i,ax in enumerate(axes.flat):\n    if i < len(wrong_images):\n        ax.imshow(wrong_images[i].squeeze(),cmap=\"gray\")\n        ax.set_title(f\"Pred: {wrong_preds[i]} | True: {wrong_labels[i]}\",fontsize=8)\n    ax.axis(\"off\")\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T19:09:41.986348Z","iopub.execute_input":"2026-08-25T19:09:41.986771Z","iopub.status.idle":"2026-08-25T19:09:42.345200Z","shell.execute_reply.started":"2026-08-25T19:09:41.986741Z","shell.execute_reply":"2026-08-25T19:09:42.344450Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"As you can see some of its wrongly predicted images have very narrow gaps between pred and true, making the model struggle. Note that this version of the model was only trained on about 35000 images, using the full 60000 images training set in the official MNIST would help generalize to such miniscule patterns better.","metadata":{}},{"cell_type":"markdown","source":"## Generate Submission","metadata":{}},{"cell_type":"code","source":"predictions = []\n\nwith torch.no_grad():\n    for images in test_loader:\n        images = images.to(device)\n        probs = sum(torch.softmax(model(images),dim=1) for model in models) / len(models)\n        predictions.extend(probs.argmax(dim=1).cpu().numpy())\n\nsubmission = pd.DataFrame({\n    \"ImageId\": range(1,len(predictions)+1),\n    \"Label\": predictions\n})\n\nsubmission.to_csv(\"submission.csv\",index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-25T19:11:04.417165Z","iopub.execute_input":"2026-08-25T19:11:04.417821Z","iopub.status.idle":"2026-08-25T19:11:12.643264Z","shell.execute_reply.started":"2026-08-25T19:11:04.417788Z","shell.execute_reply":"2026-08-25T19:11:12.642615Z"}},"outputs":[],"execution_count":null}]}