Merge changes from main

environment.yml

@@ -28,6 +28,7 @@ dependencies:
   - pyyaml
   - numba
   - pydantic
+  - joblib
 
   # Development and Testing
   - pytest

leakpro.py

@@ -7,8 +7,9 @@
 
 import joblib
 import numpy as np
-import torch
 import yaml
+from torch import load, manual_seed
+from torch.utils.data import Subset
 
 import leakpro.dev_utils.train as utils
 from leakpro import shadow_model_blueprints
@@ -79,12 +80,12 @@ def generate_user_input(configs: dict, logger: logging.Logger)->None:
     train_test_dataset = prepare_train_test_datasets(n_population, configs["data"])
 
     train_loader = get_dataloader(
-        torch.utils.data.Subset(population, train_test_dataset["train_indices"]),
+        Subset(population, train_test_dataset["train_indices"]),
         batch_size=configs["train"]["batch_size"],
         shuffle=True,
     )
     test_loader = get_dataloader(
-        torch.utils.data.Subset(population, train_test_dataset["test_indices"]),
+        Subset(population, train_test_dataset["test_indices"]),
         batch_size=configs["train"]["test_batch_size"],
     )
 
@@ -96,8 +97,8 @@ def generate_user_input(configs: dict, logger: logging.Logger)->None:
 
 
     #args = "./config/adult.yaml"  # noqa: ERA001
-    # user_args = "./config/dev_config/cifar10.yaml" # noqa: ERA001
-    user_args = "./config/dev_config/cinic10.yaml" # noqa: ERA001
+    user_args = "./config/dev_config/cifar10.yaml" # noqa: ERA001
+    #user_args = "./config/dev_config/cinic10.yaml" # noqa: ERA001
 
     with open(user_args, "rb") as f:
         user_configs = yaml.safe_load(f)
@@ -116,7 +117,7 @@ def generate_user_input(configs: dict, logger: logging.Logger)->None:
         configs = yaml.safe_load(f)
 
     # Set the random seed, log_dir and inference_game
-    torch.manual_seed(configs["audit"]["random_seed"])
+    manual_seed(configs["audit"]["random_seed"])
     np.random.seed(configs["audit"]["random_seed"])
     random.seed(configs["audit"]["random_seed"])
 
@@ -140,7 +141,7 @@ def generate_user_input(configs: dict, logger: logging.Logger)->None:
     # Load the target model parameters into the blueprint
     with open(configs["target"]["trained_model_path"], "rb") as f:
         target_model = target_model_blueprint(**target_model_metadata["model_metadata"]["init_params"])
-        target_model.load_state_dict(torch.load(f))
+        target_model.load_state_dict(load(f))
         logger.info(f"Loaded target model from {configs['target']['trained_model_path']}")
 
     # Get the population dataset

leakpro/attacks/utils/shadow_model_handler.py

@@ -66,6 +66,7 @@ def __init__(self:Self, target_model:Module, target_config:dict, config:dict, lo
             logger (logging.Logger): The logger object for logging.
 
         """
+        config = config or {}
         module_path = config.get("module_path")
         model_class_path =  config.get("model_class_path")
 

leakpro/metrics/attack_result.py

@@ -111,12 +111,6 @@ def __init__(  # noqa: PLR0913
             predicted_labels[:, true_labels == 1], axis=1
         )
 
-        sorted_indices = np.argsort(self.fp)
-        self.fp = self.fp[sorted_indices]
-        self.tp = self.tp[sorted_indices]
-        self.tn = self.tn[sorted_indices]
-        self.fn = self.fn[sorted_indices]
-
         self.roc_auc = auc(
             self.fp / (np.sum(true_labels == 0)), self.tp / (np.sum(true_labels == 1))
         )

leakpro/shadow_model_blueprints.py

@@ -1,13 +1,13 @@
 """Models for the datasets."""
-import torch
-import torch.nn.functional as F  # noqa: N812
-from torch import nn
+
+from torch import Tensor, flatten, nn
+from torch.nn import Module, functional
 from torchvision.models import resnet18
 
 from leakpro.import_helper import Self
 
 
-class NN(nn.Module):
+class NN(Module):
     """NN for Adult dataset."""
 
     def __init__(self:Self, in_shape:int, num_classes:int=10) -> None:
@@ -31,14 +31,14 @@ def __init__(self:Self, in_shape:int, num_classes:int=10) -> None:
         self.fc2 = nn.Linear(100, 50)
         self.fc3 = nn.Linear(50, num_classes)
 
-    def forward(self:Self, inputs:torch.Tensor) -> torch.Tensor:
+    def forward(self:Self, inputs:Tensor) -> Tensor:
         """Forward pass of the model."""
         inputs = inputs.flatten(1)
-        outputs = F.relu(self.fc1(inputs))
-        outputs = F.relu(self.fc2(outputs))
-        return F.relu(self.fc3(outputs))
+        outputs = functional.relu(self.fc1(inputs))
+        outputs = functional.relu(self.fc2(outputs))
+        return functional.relu(self.fc3(outputs))
 
-class ConvNet(nn.Module):
+class ConvNet(Module):
     """Convolutional Neural Network model."""
 
     def __init__(self:Self) -> None:
@@ -52,7 +52,7 @@ def __init__(self:Self) -> None:
         self.fc2 = nn.Linear(120, 84)
         self.fc3 = nn.Linear(84, 10)
 
-    def forward(self:Self, x:torch.Tensor) -> torch.Tensor:
+    def forward(self:Self, x:Tensor) -> Tensor:
         """Forward pass of the model.
 
         Args:
@@ -64,11 +64,11 @@ def forward(self:Self, x:torch.Tensor) -> torch.Tensor:
             torch.Tensor: The output tensor.
 
         """
-        x = self.pool(F.relu(self.conv1(x)))
-        x = self.pool(F.relu(self.conv2(x)))
-        x = torch.flatten(x, 1) # flatten all dimensions except batch
-        x = F.relu(self.fc1(x))
-        x = F.relu(self.fc2(x))
+        x = self.pool(functional.relu(self.conv1(x)))
+        x = self.pool(functional.relu(self.conv2(x)))
+        x = flatten(x, 1) # flatten all dimensions except batch
+        x = functional.relu(self.fc1(x))
+        x = functional.relu(self.fc2(x))
         return self.fc3(x)
 
 
@@ -88,12 +88,12 @@ def __init__(self:Self) -> None:
         self.fc2 = nn.Linear(120, 84)
         self.fc3 = nn.Linear(84, 10)
 
-    def forward(self:Self, x:torch.Tensor) -> torch.Tensor:
+    def forward(self:Self, x:Tensor) -> Tensor:
         """Forward pass of the model."""
-        x = self.pool(F.relu(self.conv1(x)))
-        x = torch.flatten(x, 1)  # flatten all dimensions except the batch
-        x = F.relu(self.fc1(x))
-        x = F.relu(self.fc2(x))
+        x = self.pool(functional.relu(self.conv1(x)))
+        x = flatten(x, 1)  # flatten all dimensions except the batch
+        x = functional.relu(self.fc1(x))
+        x = functional.relu(self.fc2(x))
         return self.fc3(x)
 
 class ResNet18(nn.Module):