update entropy TA plot

Author: michaelharborg

exp5_SCRUB_and_TA/experiment_run.py

@@ -17,7 +17,7 @@
 from src.evaluation.membership_inference_attack import MIA
 from src.evaluation.unlearning_evaluator import UnlearningEvaluator
 from prepare_image_data import get_image_unlearn_data
-from prepare_image_data_v2 import get_image_unlearn_data as get_image_unlearn_data_tsne_box
+from prepare_image_data_tsne import get_image_unlearn_data as get_image_unlearn_data_tsne_box
 import time
 import json
 

exp5_SCRUB_and_TA/mnist_config.yaml

@@ -27,11 +27,15 @@ data:
   n_channels: 3
   # cifar parameters
   split_type: "tsne_box" # "tsne_box" or "random"
-  tsne_box_coordinates: 
+  tsne_box_coordinates:
+    # x_min: -4.2 
+    # x_max: -2.6
+    # y_min: -6.25
+    # y_max: -4.8 
     # x_min: 10.2 #-6.2 
     # x_max: 11.0 # -3.4
     # y_min: -4.3 # -1.6
-    # y_max: -0.4 # 1.011
+    # y_max: -0.4 # 1.0
     x_min: -6.2 #10.2
     x_max: -3.4 #11.0
     y_min: -1.6 # -4.3
@@ -64,7 +68,7 @@ unlearn:
   method: "scrub" # "scrub" or "teacher_ascend"
   teacher_ascend:
     versions: [["ce", false], ["entropy", false], ["ce-retain", false], ["entropy-retain", false], ["ce", true], ["entropy", true], ["ce-retain", true], ["entropy-retain", true], ["ce-retain-no-reg", false], ["entropy-retain-no-reg", false]]
-    n_epochs: 150
+    n_epochs: 100
     _lambda: 2 # 64
   scrub:
     alpha: 2

exp5_SCRUB_and_TA/plot_mnist_samples.py

@@ -6,7 +6,7 @@
 # Add the parent directory to the Python path to allow for package imports
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
-from exp5_SCRUB_and_TA.prepare_image_data_v2 import download_dataset
+from exp5_SCRUB_and_TA.prepare_image_data_tsne import download_dataset
 
 def plot_mnist_samples():
     """

…p5_SCRUB_and_TA/prepare_image_data_v2.py …_SCRUB_and_TA/prepare_image_data_tsne.pyexp5_SCRUB_and_TA/prepare_image_data_v2.py renamed to exp5_SCRUB_and_TA/prepare_image_data_tsne.py exp5_SCRUB_and_TA/prepare_image_data_v2.py renamed to exp5_SCRUB_and_TA/prepare_image_data_tsne.py

@@ -198,7 +198,7 @@ def split_data_by_tsne_box(boundary: dict = None,
         return train_dataset, retain_dataset, forget_dataset, validation_dataset, forget_indices
 
 
-def get_image_unlearn_data(root_dir: str, batch_size: int, seed: int, boundary: dict = None, subsample_size: int = None):
+def get_image_unlearn_data(root_dir: str, batch_size: int, seed: int, boundary: dict = None, subsample_size: int = None, return_tsne_results: bool = False):
 
     if boundary is None:
         boundary = {
@@ -211,18 +211,31 @@ def get_image_unlearn_data(root_dir: str, batch_size: int, seed: int, boundary:
     # Make a folder if it does not exist
     coords_string = f"{boundary['x_min']}_{boundary['x_max']}_{boundary['y_min']}_{boundary['y_max']}"
     folder_name = os.path.join(root_dir, f"tsne_box_{coords_string}_{subsample_size}")
-
+    # pdb.set_trace()
     # load data if it exists
     if os.path.exists(os.path.join(folder_name, "data.pt")):
         print(f"Loading data from {folder_name}")
         data = torch.load(os.path.join(folder_name, "data.pt"), weights_only=False)
+        
+        if return_tsne_results:
+            return data['train_loader'], data['retain_loader'], data['forget_loader'], data['validation_loader'], data['forget_indices'], data['tsne_results']
+        
         return data['train_loader'], data['retain_loader'], data['forget_loader'], data['validation_loader'], data['forget_indices']
 
-    train_dataset, retain_dataset, forget_dataset, test_dataset, forget_indices = split_data_by_tsne_box(
-        boundary=boundary,
-        return_tsne_results=False,
-        subsample_size=subsample_size
-    )
+    if return_tsne_results:
+        train_dataset, retain_dataset, forget_dataset, test_dataset, forget_indices, tsne_results = split_data_by_tsne_box(
+            boundary=boundary,
+            return_tsne_results=return_tsne_results,
+            subsample_size=subsample_size
+        )
+
+    else:
+        train_dataset, retain_dataset, forget_dataset, test_dataset, forget_indices = split_data_by_tsne_box(
+            boundary=boundary,
+            return_tsne_results=return_tsne_results,
+            subsample_size=subsample_size
+        )
+    
 
     from prepare_image_data import create_image_dataloaders
 
@@ -245,9 +258,12 @@ def get_image_unlearn_data(root_dir: str, batch_size: int, seed: int, boundary:
         'retain_loader': retain_loader,
         'forget_loader': forget_loader,
         'validation_loader': validation_loader,
-        'forget_indices': forget_indices
+        'forget_indices': forget_indices,
+        'tsne_results': tsne_results
     }, os.path.join(folder_name, "data.pt"))
 
+    if return_tsne_results:
+        return train_loader, retain_loader, forget_loader, validation_loader, forget_indices, tsne_results
 
     return train_loader, retain_loader, forget_loader, validation_loader, forget_indices
 

exp5_SCRUB_and_TA/ta_scrub_experiment.py

@@ -18,6 +18,9 @@
 import graphviz
 import torch.nn as nn
 import matplotlib.pyplot as plt
+from matplotlib.colors import ListedColormap
+from scipy.interpolate import griddata
+
 from torch.utils.data import RandomSampler, SequentialSampler
 
 from src.models.neural_network import NeuralNetRS
@@ -26,9 +29,11 @@
 from src.evaluation.membership_inference_attack import MIA
 from src.evaluation.unlearning_evaluator import UnlearningEvaluator
 from prepare_image_data import get_image_unlearn_data
-from prepare_image_data_v2 import get_image_unlearn_data as get_image_unlearn_data_tsne_box
+from prepare_image_data_tsne import get_image_unlearn_data as get_image_unlearn_data_tsne_box
 import time
 import json
+import numpy as np
+
 
 def has_shuffle(dataloader) -> bool:
     """
@@ -140,6 +145,156 @@ def calculate_model_metrics(model, dataloaders, device, model_name=None, save_pa
 
     return metrics
 
+def get_nearest_neighbor_idxs(X, y, tsne_results, num_neighbors: int = 5):
+    """
+    A function that finds the nearest neighbors in t-SNE space to each point in the dataset X.
+    The neighbors are constrained such that they must belong to the same class as the sample.
+
+    Returns:
+        A numpy array of size [N x num_neighbors] containing the indices in X which have
+    """
+    # calculate distance from every point to every other point in t-SNE space
+    distances = np.linalg.norm((tsne_results[:, np.newaxis, :] - tsne_results[np.newaxis, :, :]), ord=None, axis=2)
+    np.fill_diagonal(distances, np.inf) # make sure to exclude diagonal entries
+    sorted_idxs = np.argsort(distances, axis=1) # sort distances in ascending order and get indices.
+
+    labels = y.argmax(dim=-1) # convert from OHE to labels
+    neighbor_labels = labels[sorted_idxs] # N x N (labels of all neighbors)
+
+    same_class_mask = (neighbor_labels == labels[:, None]) # N x N mask where true means the neighbor has the same class
+    masked_sorted_idxs = np.where(same_class_mask, sorted_idxs, -1) # fill different classes with index -1
+    
+    nn_idxs = np.full((X.size(0), num_neighbors), -1)
+    for i in range(X.size(0)):
+        valid_idxs = sorted_idxs[i][masked_sorted_idxs[i] != -1]
+        nn_idxs[i, :min(num_neighbors, len(valid_idxs))] = valid_idxs[:num_neighbors]
+        
+    return nn_idxs
+
+def interpolate_points_generator(X, nn_idxs, n_interp_points: int = 100):
+    """
+        X: N x M matrix of points
+        nn_idxs: A N x n_neighbors matrix of the same order as X.
+    """
+    N = X.shape[0]
+    n_neighbors = nn_idxs.shape[1]
+    interp_range = np.arange(1/n_interp_points, 1, step=1/n_interp_points)
+    # X_interp = np.zeros((N * n_interp_points * len(interp_range), X.shape[1]))
+    
+    idx = 0 
+    for i in range(X.shape[0]):
+        for j in range(n_neighbors):
+            x = X[i]
+            x_nbr = X[nn_idxs[i][j]]
+            for alpha in interp_range:
+                x_interp = x * alpha + (1-alpha) * x_nbr
+                yield x_interp.unsqueeze(0)
+
+
+def interpolate_points_batch(X, nn_idxs, n_interp_points: int = 100, batch_size: int = 512):
+    from itertools import islice
+    generator = interpolate_points_generator(X, nn_idxs, n_interp_points)
+    while True:
+        batch = list(islice(generator, batch_size))
+        if not batch:
+            break
+        yield torch.cat(batch)
+
+def create_entropy_tsne_plot(model, y, 
+                             image_generator, 
+                             tsne_results,
+                             tsne_point_generator, 
+                             tsne_bounds_x: tuple[float, float], 
+                             tsne_bounds_y: tuple[float, float],
+                             n_imgs_per_batch: int = 10, 
+                             grid_size: int = 500,
+                             device='cuda'):
+    
+    # Setup
+    class_colors = ['#003f5c', '#2f4b7c', '#665191', '#a05195', '#d45087', 
+                    '#f95d6a', '#ff7c43', '#ffa600', '#aa5382', '#eea152']
+    plt.style.use('seaborn-v0_8-paper') 
+    plt.figure(figsize=(8, 6))
+    
+    y = y.argmax(dim=-1)  # from one-hot to label
+
+    # Accumulate data
+    all_tsne = []
+    all_entropy = []
+    all_labels = []
+
+    itt = 0
+    while True:
+        imgs = next(iter(image_generator), None)
+        tsne_points = next(iter(tsne_point_generator), None)
+        if imgs is None: #or tsne_points is None:
+            break
+
+        imgs = imgs.to(device)
+        ys = y[itt * n_imgs_per_batch : (itt + 1) * n_imgs_per_batch]
+        ys = ys.repeat_interleave(imgs.size(0) // n_imgs_per_batch)
+
+        probs = model.inference(imgs)['probabilities']
+        entropies = -(torch.log(probs + 1e-8) * probs).sum(dim=-1)
+
+        all_tsne.append(tsne_points.cpu().numpy())
+        all_entropy.append(entropies.cpu().numpy())
+        all_labels.append(ys.cpu().numpy())
+
+        itt += 1
+
+    all_tsne = np.concatenate(all_tsne, axis=0)
+    all_entropy = np.concatenate(all_entropy, axis=0)
+    all_labels = np.concatenate(all_labels, axis=0)
+
+    # Interpolate entropy
+    x_min, x_max = all_tsne[:, 0].min(), all_tsne[:, 0].max()
+    y_min, y_max = all_tsne[:, 1].min(), all_tsne[:, 1].max()
+
+    grid_x, grid_y = np.meshgrid(
+        np.linspace(x_min, x_max, grid_size),
+        np.linspace(y_min, y_max, grid_size)
+    )
+
+    grid_entropy = griddata(
+        all_tsne, all_entropy, 
+        (grid_x, grid_y), 
+        #method='linear'  # This leaves NaNs in missing regions
+    )
+
+    # Set colormap with transparency for missing areas
+    cmap = plt.cm.viridis.copy()
+    cmap.set_bad(color=(1, 1, 1, 0))  # Fully transparent for NaNs
+
+    # Show entropy heatmap with transparency for missing data
+    plt.imshow(
+        grid_entropy, 
+        extent=(x_min, x_max, y_min, y_max),
+        origin='lower',
+        cmap=cmap,
+        alpha=1.0,
+        aspect='auto'
+    )
+    plt.colorbar(label='Entropy')
+
+    # Overlay t-SNE points
+    plt.scatter(tsne_results[:, 0], tsne_results[:, 1], 
+                c=y, cmap=ListedColormap(class_colors),
+                s=10, alpha=0.2)
+
+    plt.title("Interpolated Entropy in t-SNE Space")
+    plt.xlabel("t-SNE dim 1")
+    plt.ylabel("t-SNE dim 2")
+    plt.tight_layout()
+    plt.savefig('grid_entropies.png', dpi=300)
+    import pdb; pdb.set_trace()
+    plt.show()
+
+
+
+    
+    
+
 @hydra.main(config_path=".", config_name="mnist_config")
 def main(cfg):
     print("Starting experiment with configuration: %s", cfg.data.dataset_name)
@@ -172,6 +327,7 @@ def main(cfg):
         bounding_box_coords = {"x_min": cfg.data.tsne_box_coordinates.x_min, "x_max": cfg.data.tsne_box_coordinates.x_max, "y_min": cfg.data.tsne_box_coordinates.y_min, "y_max": cfg.data.tsne_box_coordinates.y_max}
         bounding_box_name = f"{cfg.data.tsne_box_coordinates.x_min}_{cfg.data.tsne_box_coordinates.x_max}_{cfg.data.tsne_box_coordinates.y_min}_{cfg.data.tsne_box_coordinates.y_max}"
         results_folder_name = f"{cfg.data.split_type}_{bounding_box_name}"
+    
     elif cfg.data.split_type == "random":
         results_folder_name = f"{cfg.data.split_type}_{cfg.data.n_forget_points}"
 
@@ -185,25 +341,28 @@ def main(cfg):
         # --------- random sampled forget set ---------
         if cfg.data.split_type == 'random':
             dataloader_train, dataloader_retain, dataloader_forget, dataloader_val, forget_idxs = get_image_unlearn_data(root_dir=dataset_dir,
-                                                                                                                        dataset_name=cfg['data']['dataset_name'],
-                                                                                                                        n_forget_points=cfg.data.n_forget_points,
-                                                                                                                        subsample_size=cfg.data.subsample_size,
-                                                                                                                        patch_size=cfg.data.patch_size,
-                                                                                                                        batch_size=cfg.data.batch_size,
-                                                                                                                        seed=cfg.data.seed)
+                                                                                                                         dataset_name=cfg['data']['dataset_name'],
+                                                                                                                         n_forget_points=cfg.data.n_forget_points,
+                                                                                                                         subsample_size=cfg.data.subsample_size,
+                                                                                                                         patch_size=cfg.data.patch_size,
+                                                                                                                         batch_size=cfg.data.batch_size,
+                                                                                                                         seed=cfg.data.seed)
         # --------- tsne box forget set ---------
         elif cfg.data.split_type == 'tsne_box':
-            dataloader_train, dataloader_retain, dataloader_forget, dataloader_val, forget_idxs = get_image_unlearn_data_tsne_box(root_dir=dataset_dir,
-                                                                                                                 batch_size=cfg.data.batch_size,
-                                                                                                                 subsample_size=cfg.data.subsample_size,
-                                                                                                                 seed=cfg.data.seed,
-                                                                                                                 boundary=bounding_box_coords)                                                               
+            dataloader_train, dataloader_retain, dataloader_forget, dataloader_val, forget_idxs, tsne_results = get_image_unlearn_data_tsne_box(root_dir=dataset_dir,
+                                                                                                                                  batch_size=cfg.data.batch_size,
+                                                                                                                                  subsample_size=cfg.data.subsample_size,
+                                                                                                                                  seed=cfg.data.seed,
+                                                                                                                                  boundary=bounding_box_coords,
+                                                                                                                                  return_tsne_results=True)
+                                                                       
         print("MNIST dataloaders created. Train size: %d, Forget size: %d", 
                    len(dataloader_train.dataset), len(dataloader_forget.dataset))
 
     else:
         raise NotImplementedError(f"The dataset {cfg.data.dataset_name} is not supported!")
 
+    
     assert (dataloader_train.dataset.y[forget_idxs] == dataloader_forget.dataset.y).all(), 'Forget indices applied to train do not correspond to the forget data!'
 
     # ========================== Initialize logger ==========================    
@@ -260,7 +419,7 @@ def main(cfg):
         else:
             # Load original model weights
             print("Loading pre-trained original model weights")
-            original_model.load_state_dict(torch.load(orig_model_path))
+            original_model.load_state_dict(torch.load(orig_model_path, map_location=DEVICE))
 
         # Calculate and save metrics for original model
         calculate_model_metrics(original_model, 
@@ -310,7 +469,7 @@ def main(cfg):
 
         else:
             print("Loading pre-trained retrained model weights")
-            retrained_model_sd = torch.load(f'{weights_dir}/retrained_model_weights_{seed}.pt')
+            retrained_model_sd = torch.load(f'{weights_dir}/retrained_model_weights_{seed}.pt', map_location=DEVICE)
             retrained_model.load_state_dict(retrained_model_sd)
 
         # Calculate and save metrics for retrained model
@@ -319,6 +478,23 @@ def main(cfg):
                                 DEVICE, 'Retrained model',
                                 save_path=f'{results_dir}/retrained_model_metrics_{seed}.json')
 
+        num_neighbors = 3
+        n_interp_points = 25
+        n_imgs_per_batch = 10
+        batch_size = n_interp_points * num_neighbors * n_imgs_per_batch
+        
+        nn_idxs = get_nearest_neighbor_idxs(dataloader_train.dataset.X, dataloader_train.dataset.y, tsne_results, num_neighbors)
+        interp_img_generator = interpolate_points_batch(dataloader_train.dataset.X, nn_idxs, n_interp_points=n_interp_points, batch_size=batch_size)
+        interp_tsne_generator = interpolate_points_batch(torch.tensor(tsne_results), nn_idxs, n_interp_points=n_interp_points, batch_size=batch_size)
+        
+        create_entropy_tsne_plot(original_model, dataloader_train.dataset.y,
+                                 interp_img_generator, tsne_results,
+                                 interp_tsne_generator,
+                                 (tsne_results[:, 0].min(), tsne_results[:, 0].max()),
+                                 (tsne_results[:, 1].min(), tsne_results[:, 1].max()),
+                                 n_imgs_per_batch, device='cuda')
+        
+        import pdb; pdb.set_trace()
 
         # ========================== Unlearn: Teacher Ascender ==========================
         if cfg.unlearn.method == 'teacher_ascend':