111 lines
4.3 KiB
Python
111 lines
4.3 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torch.utils.data import DataLoader
|
|
import torchvision.models as models
|
|
|
|
class ZeroRetrainForgettingEvaluator:
|
|
def __init__(self, unlearned_model: nn.Module, num_classes: int):
|
|
"""
|
|
Initializes the ZRF Evaluator.
|
|
|
|
Args:
|
|
unlearned_model (nn.Module): Your fine-tuned & unlearned ResNet-50.
|
|
num_classes (int): Number of classes used in your CelebA task.
|
|
"""
|
|
# select device
|
|
if torch.cuda.is_available():
|
|
self.device = torch.device("cuda")
|
|
elif hasattr(torch, "xpu") and torch.xpu.is_available():
|
|
self.device = torch.device("xpu") # For Intel GPUs using IPEX
|
|
else:
|
|
self.device = torch.device("cpu")
|
|
|
|
print(f"[INFO] Using device: {self.device}")
|
|
|
|
# prepare the unlearned model
|
|
self.unlearned_model = unlearned_model.to(self.device)
|
|
self.unlearned_model.eval()
|
|
|
|
# Instantiate a structurally matching, completely random model
|
|
print(f"[INFO] Initializing random baseline ResNet-50 with {num_classes} classes...")
|
|
self.random_model = self.get_random_model(num_classes)
|
|
self.random_model = self.random_model.to(self.device)
|
|
self.random_model.eval()
|
|
|
|
# gets randomly initialised model
|
|
# for comparison with unlearned model
|
|
def get_random_model(num_classes):
|
|
print(f"[INFO] Initializing random baseline ResNet-50 with {num_classes} classes...")
|
|
model = models.resnet50(weights=None)
|
|
model.fc = nn.Linear(model.fc.in_features, num_classes)
|
|
return model
|
|
|
|
|
|
# compute divergence
|
|
def _compute_js_divergence(self, p: torch.Tensor, q: torch.Tensor) -> float:
|
|
"""
|
|
Computes the Jensen-Shannon (JS) Divergence between two probability distributions.
|
|
|
|
Args:
|
|
p, q (Tensor): Tensors of shape (batch_size, num_classes) containing probabilities.
|
|
"""
|
|
# Avoid log(0) issues by adding a tiny epsilon
|
|
eps = 1e-12
|
|
p = torch.clamp(p, eps, 1.0)
|
|
q = torch.clamp(q, eps, 1.0)
|
|
|
|
# Calculate the midpoint distribution
|
|
m = 0.5 * (p + q)
|
|
|
|
# Compute KL Divergence natively: KL(P || M) and KL(Q || M)
|
|
kl_pm = torch.sum(p * (torch.log(p) - torch.log(m)), dim=1)
|
|
kl_qm = torch.sum(q * (torch.log(q) - torch.log(m)), dim=1)
|
|
|
|
# JS Divergence is the average of both KL divergences
|
|
js_div = 0.5 * (kl_pm + kl_qm)
|
|
|
|
# Return the mean divergence across the entire batch
|
|
return js_div.mean().item()
|
|
|
|
def evaluate_forget_class(self, dataset, batch_size: int = 32) -> float:
|
|
"""
|
|
Evaluates the unlearned model against the random model using images
|
|
from the forgotten class/identity.
|
|
|
|
Args:
|
|
dataset (Dataset): A PyTorch Dataset containing images of the forget set.
|
|
batch_size (int): Batch size for evaluation.
|
|
|
|
Returns:
|
|
float: The ZRF score (JS Divergence). A lower divergence means
|
|
the unlearned model is behaving exactly like a random model.
|
|
"""
|
|
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
|
|
|
|
total_js_div = 0.0
|
|
total_samples = 0
|
|
|
|
# No gradients needed for evaluation
|
|
with torch.no_grad():
|
|
for images, _ in dataloader:
|
|
images = images.to(self.device)
|
|
batch_len = images.size(0)
|
|
|
|
# Get raw outputs (logits)
|
|
unlearned_logits = self.unlearned_model(images)
|
|
random_logits = self.random_model(images)
|
|
|
|
# Convert logits to probability distributions via Softmax
|
|
unlearned_probs = F.softmax(unlearned_logits, dim=1)
|
|
random_probs = F.softmax(random_logits, dim=1)
|
|
|
|
# Calculate JS divergence for this batch
|
|
batch_js = self._compute_js_divergence(unlearned_probs, random_probs)
|
|
|
|
# Weighted average based on batch size (handles final smaller batches perfectly)
|
|
total_js_div += batch_js * batch_len
|
|
total_samples += batch_len
|
|
|
|
final_zrf_score = total_js_div / total_samples
|
|
return final_zrf_score |