standardize parameters

dilyabareeva · Jun 27, 2024 · d71c9e2 · d71c9e2
1 parent 6b1f0c6
commit d71c9e2
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 8 deletions.
diff --git a/src/utils/datasets/transformed_datasets/label_grouping.py b/src/utils/datasets/transformed_datasets/label_grouping.py
@@ -27,6 +27,7 @@ def __init__(
             device=device,
             p=1.0,
             subset_idx=None,  # apply with certainty, to all datapoints
+            cls_idx=None,
         )
         self.n_classes = n_classes
         self.classes = list(range(n_classes))

diff --git a/src/utils/datasets/transformed_datasets/label_poisoning.py b/src/utils/datasets/transformed_datasets/label_poisoning.py
@@ -11,18 +11,14 @@ def __init__(
         dataset: torch.utils.data.Dataset,
         n_classes: int,
         subset_idx: Optional[Union[List[int], torch.Tensor]] = None,
+        cls_idx: Optional[int] = None,
         p: int = 1.0,  # TODO: decide on default value vis-à-vis subset_idx
         seed: int = 42,
         device: str = "cpu",
     ):
 
         super().__init__(
-            dataset=dataset,
-            n_classes=n_classes,
-            seed=seed,
-            device=device,
-            p=p,
-            subset_idx=subset_idx,
+            dataset=dataset, n_classes=n_classes, seed=seed, device=device, p=p, subset_idx=subset_idx, cls_idx=cls_idx
         )
         self.poisoned_labels = {}
         for idx in range(self.perturbed_indices):

diff --git a/src/utils/datasets/transformed_datasets/sample_perturbation.py b/src/utils/datasets/transformed_datasets/sample_perturbation.py
@@ -12,7 +12,8 @@ def __init__(
         self,
         dataset: torch.utils.data.Dataset,
         n_classes: int,
-        subset_idx: Optional[Union[int, List[int], torch.Tensor]] = None,
+        subset_idx: Optional[Union[List[int], torch.Tensor]] = None,
+        cls_idx: Optional[int] = None,
         p: float = 1.0,
         seed: int = 42,
         device: str = "cpu",
@@ -25,6 +26,8 @@ def __init__(
             seed=seed,
             device=device,
             p=p,
-            subset_idx=subset_idx,  # apply with certainty, to all datapoints
+            cls_idx=cls_idx,
+            subset_idx=subset_idx,
             sample_fn=sample_fn,
+            cls_idx=cls_idx,
         )