Update app.py

app.py

@@ -5,25 +5,24 @@ import torchvision.transforms as transforms
 from torchvision import models
 import torch.nn as nn
 import torch.optim as optim
-import numpy as np
 from PIL import Image
- 
+
 # CIFAR-10 labels
 cifar10_classes = ['airplane', 'automobile', 'bird', 'cat', 'deer',
                    'dog', 'frog', 'horse', 'ship', 'truck']
- 
-# Transforms
+
+# Transforms with proper normalization for 3 channels
 transform = transforms.Compose([
     transforms.Resize((32, 32)),
     transforms.ToTensor(),
-    transforms.Normalize((0.5,), (0.5,))
+    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
 ])
- 
-# Load CIFAR-10
+
+# 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)
 testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)
- 
+
 def predict(model, image_tensor):
     model.eval()
     with torch.no_grad():
@@ -34,27 +33,45 @@ def predict(model, image_tensor):
             probs = torch.zeros_like(probs)
         pred = torch.argmax(probs).item()
     return probs, pred
- 
+
 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.
+    """
     model.train()
+    # Freeze all layers except the final fully connected layer (fc)
+    for name, param in model.named_parameters():
+        if "fc" not in name:
+            param.requires_grad = False
+
+    # Set BatchNorm layers to eval mode to prevent updating running stats
     for m in model.modules():
         if isinstance(m, nn.BatchNorm2d):
             m.eval()
- 
+
     criterion = nn.CrossEntropyLoss()
-    optimizer = optim.SGD(model.parameters(), lr=learning_rate)
- 
+    # 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, torch.tensor([label_idx]))
+        loss = -criterion(output, label_tensor)  # Negative loss for unlearning
         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}")
+        
         optimizer.zero_grad()
         loss.backward()
+        # Clip gradients to avoid explosion
+        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
         optimizer.step()
- 
+
 def evaluate_model(model, testloader):
     model.eval()
     total, correct, loss_total = 0, 0, 0.0
@@ -68,70 +85,76 @@ def evaluate_model(model, testloader):
             correct += (preds == labels).sum().item()
             loss_total += loss.item() * labels.size(0)
     return round(100 * correct / total, 2), round(loss_total / total, 4)
- 
+
 def run_unlearning(index_to_unlearn, learning_rate):
-    # Load original model
+    # Set device (CPU in this example; update as needed)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Load the original pre-trained model
     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"))
+    original_model.load_state_dict(torch.load("resnet18.pth", map_location=device))
+    original_model.to(device)
     original_model.eval()
- 
-    # Duplicate model for unlearning
+
+    # Duplicate the model for 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"))
- 
-    # Get sample
+    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
+
+    # Prediction before unlearning
     probs_before, pred_before = predict(original_model, image_tensor)
     conf_before = probs_before[label_idx].item()
- 
-    # Unlearning
+
+    # Perform unlearning on the duplicated model
     unlearn(unlearned_model, image_tensor, label_idx, learning_rate)
- 
-    # Prediction after
+
+    # Prediction after unlearning
     probs_after, pred_after = predict(unlearned_model, image_tensor)
     conf_after = probs_after[label_idx].item()
- 
-    # Evaluate full test set
+
+    # Evaluate full test set performance on 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:.4f}
- 
+
 🧽 AFTER Unlearning:
 - Prediction: {cifar10_classes[pred_after]}
 - Confidence: {conf_after:.4f}
- 
+
 📉 Confidence Drop: {conf_before - conf_after:.4f}
- 
+
 🧪 Test Set Performance:
-- Original Model: {orig_acc:.2f}%
-- Unlearned Model: {unlearn_acc:.2f}% 
+- Original Model: {orig_acc:.2f}% accuracy, Loss: {orig_loss:.4f}
+- Unlearned Model: {unlearn_acc:.2f}% accuracy, Loss: {unlearn_loss:.4f}
 """
     return result
- 
-# Gradio Interface
+
+# 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",
     title="🔁 CIFAR-10 Machine Unlearning",
     description="Load a pre-trained ResNet18 and unlearn a specific index from the CIFAR-10 training set."
 )
- 
+
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()