Update app.py

app.py

@@ -8,10 +8,12 @@ import torch.optim as optim
 from PIL import Image
 
 # CIFAR-10 labels
-cifar10_classes = ['airplane', 'automobile', 'bird', 'cat', 'deer',
-                   'dog', 'frog', 'horse', 'ship', 'truck']
+cifar10_classes = [
+    'airplane', 'automobile', 'bird', 'cat', 'deer',
+    'dog', 'frog', 'horse', 'ship', 'truck'
+]
 
-# Transforms with proper normalization for 3 channels
+# Define transformations with proper normalization for 3 channels
 transform = transforms.Compose([
     transforms.Resize((32, 32)),
     transforms.ToTensor(),
@@ -19,15 +21,30 @@ transform = transforms.Compose([
 ])
 
 # Load CIFAR-10 datasets
-trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
-testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
+trainset = torchvision.datasets.CIFAR10(
+    root='./data', train=True, download=True, transform=transform
+)
+testset = torchvision.datasets.CIFAR10(
+    root='./data', train=False, download=True, transform=transform
+)
 testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)
 
 def predict(model, image_tensor):
+    """
+    Performs a forward pass through the model and computes softmax probabilities
+    and predicted class using a numerically stable approach.
+    """
     model.eval()
     with torch.no_grad():
         outputs = model(image_tensor.unsqueeze(0))
-        probs = torch.nn.functional.softmax(outputs[0], dim=0)
+        logits = outputs[0]
+        # Use a numerically stable softmax: subtract max logit
+        max_logit = logits.max()
+        stable_logits = logits - max_logit
+        exp_logits = torch.exp(stable_logits)
+        probs = exp_logits / exp_logits.sum()
+        
+        # Check for numerical issues (if probability is exactly 0 or NaN)
         if torch.isnan(probs).any():
             print("⚠️ Warning: NaN detected in prediction probabilities")
             probs = torch.zeros_like(probs)
@@ -36,8 +53,8 @@ def predict(model, image_tensor):
 
 def unlearn(model, image_tensor, label_idx, learning_rate, steps=20):
     """
-    Performs targeted unlearning by updating only the final fully connected layer
-    using negative cross-entropy loss.
+    Performs targeted unlearning by updating only the final fully connected layer.
+    The negative cross-entropy loss drives the confidence for the target class down.
     """
     model.train()
     # Freeze all layers except the final fully connected layer (fc)
@@ -45,37 +62,41 @@ def unlearn(model, image_tensor, label_idx, learning_rate, steps=20):
         if "fc" not in name:
             param.requires_grad = False
 
-    # Set BatchNorm layers to eval mode to prevent updating running stats
+    # Set BatchNorm layers to evaluation mode to avoid updating running stats
     for m in model.modules():
         if isinstance(m, nn.BatchNorm2d):
             m.eval()
 
     criterion = nn.CrossEntropyLoss()
-    # Use Adam optimizer for parameters that require gradients (i.e. only the fc layer)
     optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=learning_rate)
-
-    # Ensure label tensor is on the same device as the image_tensor
+    
     device = image_tensor.device
     label_tensor = torch.tensor([label_idx], device=device)
     
     for i in range(steps):
         output = model(image_tensor.unsqueeze(0))
-        loss = -criterion(output, label_tensor)  # Negative loss for unlearning
+        # Negative loss to reduce confidence on the target label
+        loss = -criterion(output, label_tensor)
+        
         if torch.isnan(loss):
             print(f"❌ NaN detected in loss at step {i}. Stopping unlearning.")
             break
-        print(f"🧠 Step {i+1}/{steps} - Unlearning Loss: {loss.item():.4f}")
         
+        print(f"🧠 Step {i+1}/{steps} - Unlearning Loss: {loss.item():.4f}")
         optimizer.zero_grad()
         loss.backward()
-        # Clip gradients to avoid explosion
+        # Clip gradients to maintain stability
         torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
         optimizer.step()
 
 def evaluate_model(model, testloader):
+    """
+    Evaluates the model's accuracy and average loss on the test set.
+    """
     model.eval()
     total, correct, loss_total = 0, 0, 0.0
     criterion = nn.CrossEntropyLoss()
+    
     with torch.no_grad():
         for images, labels in testloader:
             outputs = model(images)
@@ -84,72 +105,76 @@ def evaluate_model(model, testloader):
             total += labels.size(0)
             correct += (preds == labels).sum().item()
             loss_total += loss.item() * labels.size(0)
-    return round(100 * correct / total, 2), round(loss_total / total, 4)
+    
+    accuracy = round(100 * correct / total, 2)
+    avg_loss = round(loss_total / total, 4)
+    return accuracy, avg_loss
 
 def run_unlearning(index_to_unlearn, learning_rate):
-    # Set device (CPU in this example; update as needed)
+    """
+    Loads a pre-trained ResNet18 model, performs unlearning on a single training example,
+    and compares model performance before and after unlearning.
+    """
+    # Set device (GPU if available, else CPU)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    # Load the original pre-trained model
+    
+    # Load the original pre-trained model and adjust for 10 classes
     original_model = models.resnet18(weights=None)
     original_model.fc = nn.Linear(original_model.fc.in_features, 10)
     original_model.load_state_dict(torch.load("resnet18.pth", map_location=device))
     original_model.to(device)
     original_model.eval()
-
-    # Duplicate the model for unlearning experiment
+    
+    # Duplicate the model for the unlearning experiment
     unlearned_model = models.resnet18(weights=None)
     unlearned_model.fc = nn.Linear(unlearned_model.fc.in_features, 10)
     unlearned_model.load_state_dict(torch.load("resnet18.pth", map_location=device))
     unlearned_model.to(device)
-
+    
     # Get the sample to unlearn from the training set
     image_tensor, label_idx = trainset[index_to_unlearn]
     image_tensor = image_tensor.to(device)
     label_name = cifar10_classes[label_idx]
     print(f"🗂️ Actual Label Index: {label_idx} | Label Name: {label_name}")
-
+    
     # Prediction before unlearning
     probs_before, pred_before = predict(original_model, image_tensor)
     conf_before = probs_before[label_idx].item()
-
-    # Perform unlearning on the duplicated model
+    
+    # Perform the unlearning process on the duplicated model
     unlearn(unlearned_model, image_tensor, label_idx, learning_rate)
-
+    
     # Prediction after unlearning
     probs_after, pred_after = predict(unlearned_model, image_tensor)
-    print(probs_after)
     conf_after = probs_after[label_idx].item()
-
-    # Evaluate full test set performance on both models
+    
+    print("Post-unlearning probabilities:", probs_after)
+    
+    # Evaluate the full test set performance for both models
     orig_acc, orig_loss = evaluate_model(original_model, testloader)
     unlearn_acc, unlearn_loss = evaluate_model(unlearned_model, testloader)
-
+    
     result = f"""
 📍 Index Unlearned: {index_to_unlearn}
 🗂️ Actual Label: {label_name} (Index: {label_idx})
-
 🔎 BEFORE Unlearning:
 - Prediction: {cifar10_classes[pred_before]}
 - Confidence: {conf_before:.10f}
-
 🧽 AFTER Unlearning:
 - Prediction: {cifar10_classes[pred_after]}
 - Confidence: {conf_after:.10f}
-
 📉 Confidence Drop: {conf_before - conf_after:.6f}
-
 🧪 Test Set Performance:
 - Original Model: {orig_acc:.2f}% accuracy, Loss: {orig_loss:.4f}
 - Unlearned Model: {unlearn_acc:.2f}% accuracy, Loss: {unlearn_loss:.4f}
 """
     return result
 
-# Gradio Interface for interactive unlearning demonstration
+# Gradio interface for interactive unlearning demonstration
 demo = gr.Interface(
     fn=run_unlearning,
     inputs=[
-        gr.Slider(0, len(trainset) - 1, step=1, label="Select Index to Unlearn"),
+        gr.Slider(0, len(trainset)-1, step=1, label="Select Index to Unlearn"),
         gr.Slider(0.0001, 0.01, step=0.0001, value=0.005, label="Learning Rate (for Unlearning)")
     ],
     outputs="text",