strategies tested

sets/Data.py

@@ -0,0 +1,184 @@
+from torchvision import datasets, transforms
+from torch.utils.data import Dataset, DataLoader, Subset
+import torch
+import numpy as np
+import os
+
+from enum import Enum, auto
+class Set_Name(Enum):
+    CELEBA = auto()
+    CASIAFACES = auto()
+
+# train set transform 
+def train_transform(res):
+    return transforms.Compose([
+        # ResNet expects 224 x 224 res 
+        # Inception expects 299 x 299
+        transforms.Resize((res, res)),
+        transforms.RandomHorizontalFlip(p=0.5),
+        transforms.ColorJitter(
+            brightness=0.2, 
+            contrast=0.2, 
+            saturation=0.1
+        ),
+        transforms.ToTensor(),
+        # normalise to 
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], 
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
+# test set transform
+def test_transform(res):
+    return transforms.Compose([
+        transforms.Resize((res, res)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], 
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
+# Load data with 'identity' as target and transform it
+def get_set(set_name:Set_Name):
+    return fetch_celeb_a() if set_name == Set_Name.CELEBA else fetch_casia_faces()
+
+def fetch_celeb_a():
+    return datasets.CelebA(
+        root='./data', 
+        split='all', 
+        target_type='identity', 
+        download=True, 
+        transform=None
+    )
+
+def fetch_casia_faces():
+    # location of the data (path relative to project root)
+    final_path = os.path.abspath("./data/casia-set")
+    
+    if not os.path.exists(final_path):
+        raise FileNotFoundError(
+            f"Unpacked dataset not found at {final_path}. "
+            "Please run Extractor.py first!"
+        )
+    
+    print(f"Loading unpacked CASIA dataset from: {final_path}")
+    
+    return datasets.ImageFolder(
+        root=final_path, 
+        transform=None
+    )
+
+
+
+def get_ids_and_counts(dataset):
+
+    target = get_target(dataset=dataset)
+    return  torch.unique(
+        input = target,
+        return_counts = True
+    )
+
+
+
+# filter selected identities from dataset
+# How many classes, how many images per class
+def select_ids( dataset, sample_size, class_size):
+    ids, counts = get_ids_and_counts(dataset = dataset)
+    eligible_mask = counts >= sample_size
+    eligible_ids = ids[eligible_mask].numpy()
+
+    if len(eligible_ids) < class_size:
+        raise ValueError(
+            f"Only found {len(eligible_ids)} identities with {sample_size}+ images."
+        )
+
+    # Randomly select identities
+    return np.random.choice(eligible_ids, class_size, replace=False)
+
+# optional function to get max amount of samples per class
+def select_top_ids(dataset, class_size):
+    ids, counts = get_ids_and_counts(dataset = dataset)
+
+    # sort by number of images (descending)
+    sorted_indices = torch.argsort(counts, descending = True)
+
+    top_ids = ids[sorted_indices][:class_size].numpy()
+
+    return np.array(top_ids, dtype=int)
+
+
+def get_target(dataset):
+    """
+    Unified target extractor. 
+    Instantly reads raw dataset arrays or safely scales down to unpack wrapped Subsets.
+    """
+    if hasattr(dataset, 'identity'):
+        # celebA
+        targets = dataset.identity
+    elif hasattr(dataset, 'targets'):
+        # others
+        targets = dataset.targets
+    else:
+        # If it's an IdentitySubset or standard Subset, extract mapped targets sequentially
+        # This guarantees we get the 0 -> (n-1) remapped labels
+        targets = [dataset[i][1] for i in range(len(dataset))]
+
+    if not isinstance(targets, torch.Tensor):
+        targets = torch.tensor(targets)
+        
+    return targets
+
+
+# split class images to train and test set.
+def get_indices(dataset, identities, split_at, size = 30):
+
+    if split_at >= size: # debug safety
+        raise ValueError(f"Split point ({split_at}) must be less than total size ({size}).")
+    
+    train_indices = []
+    test_indices = []
+
+    target = get_target(dataset=dataset)
+
+    #training_sample = int(sample_size * training_ratio)
+    np.random.seed(42)
+    for person_id in identities:
+        # Get all indices for this specific person
+        indices = torch.where(target == person_id)[0].numpy()
+
+        # Shuffle the indices for this person
+        np.random.shuffle(indices)
+        
+        # split data to testing and training
+        train_indices.extend(indices[:split_at])
+        test_indices.extend(indices[split_at:size])
+
+    return train_indices, test_indices
+
+
+
+def get_unlearning_loaders(dataset: Dataset, forget_class_idx: int, batch_size: int = 32) -> tuple[DataLoader, DataLoader]:
+    """
+    Splits an IdentitySubset or standard Dataset into forget and retain sets 
+    based on a remapped target class index.
+    """
+    # extract targets
+    targets = get_target(dataset=dataset)
+        
+    # mask indices local to this subset
+    forget_indices = torch.where(targets == forget_class_idx)[0].tolist()
+    retain_indices = torch.where(targets != forget_class_idx)[0].tolist()
+    
+    # PyTorch Subsets
+    forget_subset = Subset(dataset, forget_indices)
+    retain_subset = Subset(dataset, retain_indices)
+    
+    # DataLoaders
+    forget_loader = DataLoader(forget_subset, batch_size=batch_size, shuffle=False)
+    retain_loader = DataLoader(retain_subset, batch_size=batch_size, shuffle=True)
+    
+    print(f"[Data Split] Local Class {forget_class_idx}: {len(forget_subset)} samples | Remaining Classes: {len(retain_subset)} samples.")
+    
+    return forget_loader, retain_loader