Create app.py

app.py

@@ -0,0 +1,128 @@
+import gradio as gr
+import torch
+import torchvision
+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
+transform = transforms.Compose([
+    transforms.Resize((32, 32)),
+    transforms.ToTensor(),
+    transforms.Normalize((0.5,), (0.5,))
+])
+ 
+# Load CIFAR-10
+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():
+        outputs = model(image_tensor.unsqueeze(0))
+        probs = torch.nn.functional.softmax(outputs[0], dim=0)
+        pred = torch.argmax(probs).item()
+    return probs, pred
+ 
+def unlearn(model, image_tensor, label_idx, learning_rate, steps=10):
+    model.train()
+    for m in model.modules():
+        if isinstance(m, nn.BatchNorm2d):
+            m.eval()
+ 
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.SGD(model.parameters(), lr=learning_rate)
+ 
+    for _ in range(steps):
+        output = model(image_tensor.unsqueeze(0))
+        loss = -criterion(output, torch.tensor([label_idx]))
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+ 
+def evaluate_model(model, testloader):
+    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)
+            _, preds = torch.max(outputs, 1)
+            loss = criterion(outputs, labels)
+            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)
+ 
+def run_unlearning(index_to_unlearn, learning_rate):
+    # Load original 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.eval()
+ 
+    # Duplicate model for unlearning
+    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
+    image_tensor, label_idx = trainset[index_to_unlearn]
+    label_name = cifar10_classes[label_idx]
+ 
+    # Prediction before
+    probs_before, pred_before = predict(original_model, image_tensor)
+    conf_before = probs_before[label_idx].item()
+ 
+    # Unlearning
+    unlearn(unlearned_model, image_tensor, label_idx, learning_rate)
+ 
+    # Prediction after
+    probs_after, pred_after = predict(unlearned_model, image_tensor)
+    conf_after = probs_after[label_idx].item()
+ 
+    # Evaluate full test set
+    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} | Label: {label_name}
+ 
+🔎 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}% | Loss: {orig_loss:.4f}
+- Unlearned Model: {unlearn_acc:.2f}% | Loss: {unlearn_loss:.4f}
+"""
+    return result
+ 
+# Gradio Interface
+demo = gr.Interface(
+    fn=run_unlearning,
+    inputs=[
+        gr.Slider(0, len(trainset)-1, step=1, label="Select Index to Unlearn"),
+        gr.Slider(0.0001, 0.1, step=0.0001, value=0.01, 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()