helper code to visualize all mnist test data at once, instead of in iterative batches

Author: vyinkabanjo

lab2/solutions/PT_Part1_MNIST_Solution.ipynb

@@ -861,7 +861,7 @@
       "outputs": [],
       "source": [
         "print(\"Label of this digit is:\", test_label)\n",
-        "plt.imshow(test_image[0,0,:,:].cpu(), cmap=plt.cm.binary)\n",
+        "plt.imshow(test_image[0,0,:,:], cmap=plt.cm.binary)\n",
         "comet_model_2.log_figure(figure=plt)"
       ]
     },
@@ -871,7 +871,45 @@
         "id": "ygh2yYC972ne"
       },
       "source": [
-        "It is! Let's visualize the classification results on the MNIST dataset. We will plot images from the test dataset along with their predicted label, as well as a histogram that provides the prediction probabilities for each of the digits:"
+        "It is! Let's visualize the classification results on the MNIST dataset. We will plot images from the test dataset along with their predicted label, as well as a histogram that provides the prediction probabilities for each of the digits.\n",
+        "\n",
+        "Recall that in PyTorch the MNIST dataset is typically accessed using a DataLoader to iterate through the test set in smaller, manageable batches. By appending the predictions, test labels, and test images from each batch, we will first gradually accumulate all the data needed for visualization into singular variables to observe our model's predictions."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Initialize variables to store all data\n",
+        "all_predictions = []\n",
+        "all_labels = []\n",
+        "all_images = []\n",
+        "\n",
+        "# Process test set in batches\n",
+        "with torch.no_grad():\n",
+        "    for images, labels in testset_loader:\n",
+        "        outputs = cnn_model(images)\n",
+        "\n",
+        "         # Apply softmax to get probabilities\n",
+        "        probabilities = torch.nn.functional.softmax(outputs, dim=1)\n",
+        "\n",
+        "        # Get predicted classes\n",
+        "        _, predicted = torch.max(outputs, 1)\n",
+        "\n",
+        "        all_predictions.append(probabilities) \n",
+        "        all_labels.append(labels)\n",
+        "        all_images.append(images)\n",
+        "\n",
+        "all_predictions = torch.cat(all_predictions)  # Shape: (total_samples, num_classes)\n",
+        "all_labels = torch.cat(all_labels)            # Shape: (total_samples,)\n",
+        "all_images = torch.cat(all_images)            # Shape: (total_samples, 1, 28, 28)\n",
+        "\n",
+        "# Convert tensors to NumPy for compatibility with plotting functions\n",
+        "predictions = all_predictions.numpy()  # Shape: (total_samples, num_classes)\n",
+        "test_labels = all_labels.numpy()       # Shape: (total_samples,)\n",
+        "test_images = all_images.numpy()       # Shape: (total_samples, 1, 28, 28)"
       ]
     },
     {