fix: randomize all layer types

quanda/benchmarks/heuristics/mixed_datasets.py

@@ -346,7 +346,7 @@ def download(cls, name: str, cache_dir: str, device: str, *args, **kwargs):
             adversarial_dir=adversarial_dir,
             adversarial_label=bench_state["adversarial_label"],
             adversarial_transform=adversarial_transform,
-            adv_train_indices = adv_train_indices,
+            adv_train_indices=adv_train_indices,
             data_transform=dataset_transform,
             checkpoint_paths=checkpoint_paths,
         )

quanda/metrics/heuristics/model_randomization.py

@@ -193,19 +193,35 @@ def load_state_dict(self, state_dict: dict):
         self.results = state_dict["results_dict"]
         self.rand_model.load_state_dict(state_dict["rnd_model"])
 
+    def _randomize_parameter(self, param, parent, param_name):
+        """Reset or randomize a parameter.
+
+        Parameters
+        ----------
+        param : torch.Tensor
+            The parameter tensor.
+        parent : torch.nn.Module
+            The parent module of the parameter.
+        param_name : str
+            The name of the parameter.
+
+        """
+        if hasattr(parent, "reset_parameters"):
+            torch.manual_seed(self.seed)
+            parent.reset_parameters()
+        else:
+            torch.nn.init.normal_(param, generator=self.generator)
+            parent.__setattr__(param_name, torch.nn.Parameter(param))
+
     def _randomize_model(self) -> Tuple[torch.nn.Module, List[str]]:
         """Randomize the model parameters.
 
-        Currently, only linear and convolutional layers are supported.
-
         Returns
         -------
         torch.nn.Module
             The randomized model.
 
         """
-        # TODO: Add support for other layer types.
-
         rand_model = copy.deepcopy(self.model)
         rand_checkpoints = []
 
@@ -214,24 +230,14 @@ def _randomize_model(self) -> Tuple[torch.nn.Module, List[str]]:
 
             for name, param in list(rand_model.named_parameters()):
                 parent = get_parent_module_from_name(rand_model, name)
-                # TODO: currently only linear layer is randomized
-                if isinstance(parent, (torch.nn.Linear)):
-                    random_param_tensor = torch.nn.init.normal_(
-                        param, generator=self.generator
-                    )
-                    parent.__setattr__(
-                        name.split(".")[-1],
-                        torch.nn.Parameter(random_param_tensor),
-                    )
-
-            # save randomized checkpoint
+                param_name = name.split(".")[-1]
+                self._randomize_parameter(param, parent, param_name)
+
+            # Save randomized checkpoint
             chckpt_path = os.path.join(
                 self.cache_dir, f"{self.model_id}_rand_{i}.pth"
             )
-            torch.save(
-                rand_model.state_dict(),
-                chckpt_path,
-            )
+            torch.save(rand_model.state_dict(), chckpt_path)
             rand_checkpoints.append(chckpt_path)
 
         return rand_model, rand_checkpoints

tests/benchmarks/heuristics/test_model_randomization.py

@@ -30,7 +30,7 @@
                 "similarity_metric": cosine_similarity,
             },
             None,
-            0.5208332538604736,
+            0.717261791229248,
         ),
         (
             "mnist2",
@@ -49,7 +49,7 @@
                 "similarity_metric": cosine_similarity,
             },
             None,
-            0.5208332538604736,
+            0.717261791229248,
         ),
     ],
 )
@@ -120,7 +120,7 @@ def test_model_randomization(
                 "similarity_metric": cosine_similarity,
                 "load_from_disk": True,
             },
-            0.509926438331604,
+            0.19356615841388702,
         ),
     ],
 )

tests/conftest.py

@@ -7,6 +7,7 @@
 import torch
 import torchvision
 from torch.utils.data import TensorDataset
+from torchvision.models import vit_b_16, resnet18
 
 from quanda.benchmarks.downstream_eval import (
     ClassDetection,
@@ -123,6 +124,20 @@ def load_mnist_model():
     return model
 
 
+@pytest.fixture
+def load_mnist_model_with_custom_param():
+    """Load a pre-trained LeNet classification model with a custom parameter
+    (architecture at quantus/helpers/models)."""
+    model = LeNet()
+    model.load_state_dict(
+        torch.load(
+            "tests/assets/mnist", map_location="cpu", pickle_module=pickle
+        )
+    )
+    model.custom_param = torch.nn.Parameter(torch.randn(4))
+    return model
+
+
 @pytest.fixture
 def load_mnist_last_checkpoint():
     """Load the path to the last checkpoint of a pre-trained LeNet classification model."""
@@ -455,3 +470,13 @@ def load_pretrained_models_lds():
         model.load_state_dict(torch.load(path, map_location="cpu"))
         models.append(model)
     return models
+
+
+@pytest.fixture
+def load_vit():
+    return vit_b_16()
+
+
+@pytest.fixture
+def load_resnet():
+    return resnet18()