Merge pull request #116 from aidotse/single_sm

Single softmax

config/audit.yaml

@@ -4,28 +4,28 @@ audit:  # Configurations for auditing
     rmia:
       training_data_fraction: 0.1  # Fraction of the auxilary dataset to use for this attack (in each shadow model training)
       attack_data_fraction: 0.1 # Fraction of auxiliary dataset to sample from during attack
-      num_shadow_models: 4 # Number of shadow models to train
+      num_shadow_models: 3 # Number of shadow models to train
       online: True # perform online or offline attack
       temperature: 2
       gamma: 2.0
       offline_a: 0.33 # parameter from which we compute p(x) from p_OUT(x) such that p_IN(x) = a p_OUT(x) + b.
       offline_b: 0.66
-    qmia:
-      training_data_fraction: 0.1  # Fraction of the auxilary dataset (data without train and test indices) to use for training the quantile regressor
-      epochs: 5  # Number of training epochs for quantile regression
-    population:
-      attack_data_fraction: 0.1  # Fraction of the auxilary dataset to use for this attack
-    lira:
-      training_data_fraction: 0.5  # Fraction of the auxilary dataset to use for this attack (in each shadow model training)
-      num_shadow_models: 4 # Number of shadow models to train
-      online: False # perform online or offline attack
-      memorization: True # Use memorization score to boost performance
-    loss_traj:
-      training_distill_data_fraction : 0.2 # Fraction of the auxilary dataset to use for training the distillation models D_s = (1-D_KD)/2
-      number_of_traj: 1 # Number of epochs (number of points in the loss trajectory)
-      label_only: "False" # True or False
-      attack_data_dir: "./leakpro_output/attack_objects/loss_traj"
-      mia_classifier_epochs: 10
+    # qmia:
+    #   training_data_fraction: 0.1  # Fraction of the auxilary dataset (data without train and test indices) to use for training the quantile regressor
+    #   epochs: 5  # Number of training epochs for quantile regression
+    # population:
+    #   attack_data_fraction: 0.1  # Fraction of the auxilary dataset to use for this attack
+    # lira:
+    #   training_data_fraction: 0.5  # Fraction of the auxilary dataset to use for this attack (in each shadow model training)
+    #   num_shadow_models: 4 # Number of shadow models to train
+    #   online: False # perform online or offline attack
+    #   memorization: True # Use memorization score to boost performance
+    # loss_traj:
+    #   training_distill_data_fraction : 0.2 # Fraction of the auxilary dataset to use for training the distillation models D_s = (1-D_KD)/2
+    #   number_of_traj: 1 # Number of epochs (number of points in the loss trajectory)
+    #   label_only: "False" # True or False
+    #   attack_data_dir: "./leakpro_output/attack_objects/loss_traj"
+    #   mia_classifier_epochs: 10
 
   report_log: "./leakpro_output/results"  # Folder to save the auditing report
   config_log: "./leakpro_output/config"  # Folder to save the configuration files
@@ -51,7 +51,7 @@ shadow_model:
   # Name of the class to instantiate from the specified file
   model_class: "ResNet18" #"ConvNet"
   batch_size: 256
-  epochs: 30
+  epochs: 1
   optimizer: 
     name: sgd #adam, sgd, rmsprop
     lr: 0.01

config/dev_config/cifar10.yaml

@@ -5,7 +5,7 @@ run: # Configurations for a specific run
 train: # Configuration for training
   type: pytorch # Training framework (we only support pytorch now).
   num_target_model: 1 #Integer number for indicating how many target models we want to audit for the privacy game
-  epochs: 20 # Integer number for indicating the epochs for training target model. For speedyresnet, it uses its own number of epochs.
+  epochs: 1 # Integer number for indicating the epochs for training target model. For speedyresnet, it uses its own number of epochs.
   batch_size: 256 # Integer number for indicating batch size for training the target model. For speedyresnet, it uses its own batch size.
   optimizer: SGD # String which indicates the optimizer. We support Adam and SGD. For speedyresnet, it uses its own optimizer.
   learning_rate: 0.01 # Float number for indicating learning rate for training the target model. For speedyresnet, it uses its own learning_rate.

leakpro/attacks/mia_attacks/rmia.py

@@ -4,6 +4,7 @@
 
 from leakpro.attacks.mia_attacks.abstract_mia import AbstractMIA
 from leakpro.attacks.utils.shadow_model_handler import ShadowModelHandler
+from leakpro.attacks.utils.utils import softmax_logits
 from leakpro.import_helper import Self
 from leakpro.metrics.attack_result import CombinedMetricResult
 from leakpro.signals.signal import ModelLogits
@@ -87,35 +88,6 @@ def description(self:Self) -> dict:
             "detailed": detailed_str,
         }
 
-    def softmax(self:Self, all_logits:np.ndarray,
-                true_label_indices:np.ndarray,
-                return_full_distribution:bool=False) -> np.ndarray:
-        """Compute the softmax function.
-
-        Args:
-        ----
-            all_logits (np.ndarray): Logits for each class.
-            true_label_indices (np.ndarray): Indices of the true labels.
-            return_full_distribution (bool, optional): return the full distribution or just the true class probabilities.
-
-        Returns:
-        -------
-            np.ndarray: Softmax output.
-
-        """
-        logit_signals = all_logits / self.temperature
-        max_logit_signals = np.max(logit_signals,axis=2)
-        logit_signals = logit_signals - max_logit_signals.reshape(1,-1,1)
-        exp_logit_signals = np.exp(logit_signals)
-        exp_logit_sum = np.sum(exp_logit_signals, axis=2)
-
-        if return_full_distribution is False:
-            true_exp_logit =  exp_logit_signals[:, np.arange(exp_logit_signals.shape[1]), true_label_indices]
-            output_signal = true_exp_logit / exp_logit_sum
-        else:
-            output_signal = exp_logit_signals / exp_logit_sum[:,:,np.newaxis]
-        return output_signal
-
     def prepare_attack(self:Self) -> None:
         """Prepare data needed for running the attack on the target model and dataset.
 
@@ -159,14 +131,14 @@ def prepare_attack(self:Self) -> None:
             # run points through real model to collect the logits
             logits_theta = np.array(self.signal([self.target_model], self.handler, chosen_attack_data_indices))
             # collect the softmax output of the correct class
-            p_z_given_theta = self.softmax(logits_theta, z_true_labels)
+            n_attack_points = len(chosen_attack_data_indices)
+            p_z_given_theta = softmax_logits(logits_theta, self.temperature)[:,np.arange(n_attack_points),z_true_labels]
 
             # run points through shadow models and collect the logits
             logits_shadow_models = self.signal(self.shadow_models, self.handler, chosen_attack_data_indices)
             # collect the softmax output of the correct class for each shadow model
-            p_z_given_shadow_models = [self.softmax(np.array(x).reshape(1,*x.shape), z_true_labels) for x in logits_shadow_models]
-            # stack the softmax output of the correct class for each shadow model to dimension # models x # data points
-            p_z_given_shadow_models = np.array(p_z_given_shadow_models).squeeze()
+            sm_logits_shadow_models = [softmax_logits(x, self.temperature) for x in logits_shadow_models]
+            p_z_given_shadow_models = np.array([x[np.arange(n_attack_points),z_true_labels] for x in sm_logits_shadow_models])
 
             # evaluate the marginal p(z)
             p_z = np.mean(p_z_given_shadow_models, axis=0) if len(self.shadow_models) > 1 else p_z_given_shadow_models.squeeze()
@@ -209,13 +181,14 @@ def _online_attack(self:Self) -> None:
         # run points through target model to get logits
         logits_theta = np.array(self.signal([self.target_model], self.handler, audit_data_indices))
         # collect the softmax output of the correct class
-        p_x_given_target_model = self.softmax(logits_theta, ground_truth_indices)
+        n_audit_points = len(audit_data_indices)
+        p_x_given_target_model = softmax_logits(logits_theta, self.temperature)[:,np.arange(n_audit_points),ground_truth_indices]
 
         # run points through shadow models, colelct logits and compute p(x)
         logits_shadow_models = self.signal(self.shadow_models, self.handler, audit_data_indices)
-        p_x_given_shadow_models = [self.softmax(np.array(x).reshape(1,*x.shape), ground_truth_indices)
-                                    for x in logits_shadow_models]
-        p_x_given_shadow_models = np.array(p_x_given_shadow_models).squeeze()
+        sm_shadow_models = [softmax_logits(x, self.temperature) for x in logits_shadow_models]
+        p_x_given_shadow_models = np.array([x[np.arange(n_audit_points),ground_truth_indices] for x in sm_shadow_models])
+
         p_x = np.mean(p_x_given_shadow_models, axis=0) if len(self.shadow_models) > 1 else p_x_given_shadow_models.squeeze()
         # compute the ratio of p(x|theta) to p(x)
         ratio_x = p_x_given_target_model / (p_x + self.epsilon)
@@ -242,14 +215,14 @@ def _online_attack(self:Self) -> None:
         # run points through real model to collect the logits
         logits_target_model = np.array(self.signal([self.target_model], self.handler, self.attack_data_index))
         # collect the softmax output of the correct class
-        p_z_given_target_model = self.softmax(logits_target_model, z_true_labels)
+        n_attack_points = len(self.attack_data_index)
+        p_z_given_target_model = softmax_logits(logits_target_model, self.temperature)[:,np.arange(n_attack_points),z_true_labels]
 
         # run points through shadow models and collect the logits
         logits_shadow_models = self.signal(self.shadow_models, self.handler, self.attack_data_index)
         # collect the softmax output of the correct class for each shadow model
-        p_z_given_shadow_models = [self.softmax(np.array(x).reshape(1,*x.shape), z_true_labels) for x in logits_shadow_models]
-        # stack the softmax output of the correct class for each shadow model to dimension # models x # data points
-        p_z_given_shadow_models = np.array(p_z_given_shadow_models).squeeze()
+        sm_shadow_models = [softmax_logits(x, self.temperature) for x in logits_shadow_models]
+        p_z_given_shadow_models = np.array([x[np.arange(n_attack_points),z_true_labels] for x in sm_shadow_models])
 
         # evaluate the marginal p(z) by averaging over the OUT models
         p_z = np.zeros((len(audit_data_indices), len(self.attack_data_index)))
@@ -275,18 +248,16 @@ def _offline_attack(self:Self) -> None:
         ground_truth_indices = self.handler.get_labels(self.audit_dataset["data"])
         assert np.issubdtype(ground_truth_indices.dtype, np.integer)
 
-        p_x_given_target_model = self.softmax(logits_theta, ground_truth_indices)
+        n_audit_points = len(self.audit_dataset["data"])
+        p_x_given_target_model = softmax_logits(logits_theta, self.temperature)[:,np.arange(n_audit_points),ground_truth_indices]
 
         # run points through shadow models and collect the logits
         logits_shadow_models = self.signal(self.shadow_models, self.handler, self.audit_dataset["data"])
         # collect the softmax output of the correct class for each shadow model
-        p_x_given_shadow_models = [
-            self.softmax(np.array(x).reshape(1,*x.shape), ground_truth_indices)
-            for x in logits_shadow_models
-        ]
-        # stack the softmax output of the correct class for each shadow model
-        # to dimension # models x # data points
-        p_x_given_shadow_models = np.array(p_x_given_shadow_models).squeeze()
+        # Stack to dimension # models x # data points
+        sm_shadow_models = [softmax_logits(x, self.temperature) for x in logits_shadow_models]
+        p_x_given_shadow_models = np.array([x[np.arange(n_audit_points),ground_truth_indices] for x in sm_shadow_models])
+
         # evaluate the marginal p_out(x) by averaging the output of the shadow models
         p_x_out = np.mean(p_x_given_shadow_models, axis=0) if len(self.shadow_models) > 1 else p_x_given_shadow_models.squeeze()
 

leakpro/attacks/utils/boosting.py

@@ -3,9 +3,9 @@
 import logging
 
 import numpy as np
-import torch
 from tqdm import tqdm
 
+from leakpro.attacks.utils.utils import softmax_logits
 from leakpro.import_helper import Self
 from leakpro.model import PytorchModel
 from leakpro.signals.signal import ModelLogits, ModelRescaledLogits
@@ -119,7 +119,7 @@ def _memorization_score(self:Self) -> None:
         logits_function = ModelLogits()
         logits = np.swapaxes(logits_function(self.shadow_models, self.handler, self.audit_data_indices, self.batch_size), 0, 1)
 
-        logits = self.softmax_logits(logits)
+        logits = softmax_logits(logits)
 
         if self.online:
             for i, (logit, mask, label) in tqdm(enumerate(zip(logits, self.in_indices_mask, self.audit_data_labels)),
@@ -136,7 +136,7 @@ def _memorization_score(self:Self) -> None:
             #       filtering of IN- vs. OUT-samples.
             target_logits = np.swapaxes(logits_function([self.target_model], self.handler, self.audit_data_indices,\
                                         self.batch_size), 0, 1).squeeze()
-            target_logits = self.softmax_logits(target_logits)
+            target_logits = softmax_logits(target_logits)
             for i, (logit, target_logit, label) in tqdm(enumerate(zip(logits, target_logits, self.audit_data_labels)),
                                                 total=len(logits),
                                                 desc="Calculating memorization score",
@@ -182,20 +182,6 @@ def _privacy_score(self:Self) -> None:
 
         self.privacy_score = np.asarray(privacy_score)
 
-    def softmax_logits(self:Self, logits: np.ndarray) -> np.ndarray:
-        """Rescale logits to (0, 1).
-
-        Args:
-        ----
-            logits ( len(dataset) x ... x nb_classes ): Logits to be rescaled.
-
-        """
-        logits = torch.from_numpy(logits)
-        logits = logits - torch.max(logits, dim=-1, keepdim=True).values
-        logits = torch.exp(logits)
-        logits = logits/torch.sum(logits, dim=-1, keepdim=True)
-        return logits.numpy()
-
     def adjust_memorization_mask(self:Self) -> list:
         """Adjust thesholds to achieve the desired amount or percentile of most vulnerable datapoints."""
 

leakpro/attacks/utils/utils.py

@@ -0,0 +1,22 @@
+"""Utility functions for attacks."""
+
+import numpy as np
+from torch import exp, from_numpy, max, sum
+
+
+def softmax_logits(logits: np.ndarray, temp:float=1.0, dimension:int=-1) -> np.ndarray:
+    """Rescale logits to (0, 1).
+
+    Args:
+    ----
+        logits ( len(dataset) x ... x nb_classes ): Logits to be rescaled.
+        temp (float): Temperature for softmax.
+        dimension (int): Dimension to apply softmax.
+
+    """
+    logits = from_numpy(logits) / temp
+    logits = logits - max(logits, dim=dimension, keepdim=True).values
+    logits = exp(logits)
+    logits = logits/sum(logits, dim=dimension, keepdim=True)
+    return logits.numpy()
+