Create misclassified_image.py

misclassified_image.py

@@ -0,0 +1,57 @@
+import torch
+import torchvision
+import torchvision.transforms as transforms
+
+def get_cifar10_dataloaders(batch_size=64):
+
+    transform_test = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+    ])
+
+    testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
+    testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=2)
+
+    return testloader
+
+#testloader = get_cifar10_dataloaders()
+
+def get_misclassified_data(model, device, test_loader):
+    """
+    Function to run the model on test set and return misclassified images
+    :param model: Network Architecture
+    :param device: CPU/GPU
+    :param test_loader: DataLoader for test set
+    """
+    # Prepare the model for evaluation i.e. drop the dropout layer
+    model.eval()
+
+    # List to store misclassified Images
+    misclassified_data = []
+
+    # Reset the gradients
+    with torch.no_grad():
+        # Extract images, labels in a batch
+        for data, target in test_loader:
+
+            # Migrate the data to the device
+            data, target = data.to(device), target.to(device)
+
+            # Extract single image, label from the batch
+            for image, label in zip(data, target):
+
+                # Add batch dimension to the image
+                image = image.unsqueeze(0)
+
+                # Get the model prediction on the image
+                output = model(image)
+
+                # Convert the output from one-hot encoding to a value
+                pred = output.argmax(dim=1, keepdim=True)
+
+                # If prediction is incorrect, append the data
+                if pred != label:
+                    misclassified_data.append((image, label, pred))
+    return misclassified_data
+##################
+#misclassified_data = get_misclassified_data(model,'cpu', testloader)