refactor: reorder model, normalize, device params

spencerwooo · Jan 20, 2024 · a4ba556 · a4ba556
1 parent 3a85a1e
commit a4ba556
Show file tree

Hide file tree

Showing 14 changed files with 38 additions and 38 deletions.
diff --git a/src/torchattack/admix.py b/src/torchattack/admix.py
@@ -17,6 +17,7 @@ def __init__(
         self,
         model: nn.Module,
         normalize: Callable[[torch.Tensor], torch.Tensor] | None,
+        device: torch.device | None = None,
         eps: float = 8 / 255,
         steps: int = 10,
         alpha: float | None = None,
@@ -27,13 +28,13 @@ def __init__(
         clip_min: float = 0.0,
         clip_max: float = 1.0,
         targeted: bool = False,
-        device: torch.device | None = None,
     ) -> None:
         """Initialize the Admix attack.
 
         Args:
             model: The model to attack.
             normalize: A transform to normalize images.
+            device: Device to use for tensors. Defaults to cuda if available.
             eps: The maximum perturbation. Defaults to 8/255.
             steps: Number of steps. Defaults to 10.
             alpha: Step size, `eps / steps` if None. Defaults to None.
@@ -44,10 +45,9 @@ def __init__(
             clip_min: Minimum value for clipping. Defaults to 0.0.
             clip_max: Maximum value for clipping. Defaults to 1.0.
             targeted: Targeted attack if True. Defaults to False.
-            device: Device to use for tensors. Defaults to cuda if available.
         """
 
-        super().__init__(device, normalize)
+        super().__init__(normalize, device)
 
         self.model = model
         self.eps = eps

diff --git a/src/torchattack/base.py b/src/torchattack/base.py
@@ -9,8 +9,8 @@ class Attack(ABC):
 
     def __init__(
         self,
-        device: torch.device | None,
         normalize: Callable[[torch.Tensor], torch.Tensor] | None,
+        device: torch.device | None,
     ) -> None:
         super().__init__()
 

diff --git a/src/torchattack/deepfool.py b/src/torchattack/deepfool.py
@@ -17,27 +17,27 @@ def __init__(
         self,
         model: nn.Module,
         normalize: Callable[[torch.Tensor], torch.Tensor] | None,
+        device: torch.device | None = None,
         steps: int = 100,
         overshoot: float = 0.02,
         num_classes: int = 10,
         clip_min: float = 0.0,
         clip_max: float = 1.0,
-        device: torch.device | None = None,
     ) -> None:
         """Initialize the DeepFool attack.
 
         Args:
             model: The model to attack.
             normalize: A transform to normalize images.
+            device: Device to use for tensors. Defaults to cuda if available.
             steps: Number of steps. Defaults to 100.
             overshoot: Overshoot parameter for noise control. Defaults to 0.02.
             num_classes: Number of classes to consider. Defaults to 10.
             clip_min: Minimum value for clipping. Defaults to 0.0.
             clip_max: Maximum value for clipping. Defaults to 1.0.
-            device: Device to use for tensors. Defaults to cuda if available.
         """
 
-        super().__init__(device, normalize)
+        super().__init__(normalize, device)
 
         self.model = model
         self.steps = steps

diff --git a/src/torchattack/difgsm.py b/src/torchattack/difgsm.py
@@ -18,6 +18,7 @@ def __init__(
         self,
         model: nn.Module,
         normalize: Callable[[torch.Tensor], torch.Tensor] | None,
+        device: torch.device | None = None,
         eps: float = 8 / 255,
         steps: int = 10,
         alpha: float | None = None,
@@ -27,7 +28,6 @@ def __init__(
         clip_min: float = 0.0,
         clip_max: float = 1.0,
         targeted: bool = False,
-        device: torch.device | None = None,
     ) -> None:
         """Initialize the DI-FGSM attack.
 
@@ -40,6 +40,7 @@ def __init__(
         Args:
             model: The model to attack.
             normalize: A transform to normalize images.
+            device: Device to use for tensors. Defaults to cuda if available.
             eps: The maximum perturbation. Defaults to 8/255.
             steps: Number of steps. Defaults to 10.
             alpha: Step size, `eps / steps` if None. Defaults to None.
@@ -49,10 +50,9 @@ def __init__(
             clip_min: Minimum value for clipping. Defaults to 0.0.
             clip_max: Maximum value for clipping. Defaults to 1.0.
             targeted: Targeted attack if True. Defaults to False.
-            device: Device to use for tensors. Defaults to cuda if available.
         """
 
-        super().__init__(device, normalize)
+        super().__init__(normalize, device)
 
         self.model = model
         self.eps = eps

diff --git a/src/torchattack/fgsm.py b/src/torchattack/fgsm.py
@@ -28,14 +28,14 @@ def __init__(
         Args:
             model: A torch.nn.Module network model.
             normalize: A transform to normalize images.
+            device: Device to use for tensors. Defaults to cuda if available.
             eps: Maximum perturbation measured by Linf. Defaults to 8/255.
             clip_min: Minimum value for clipping. Defaults to 0.0.
             clip_max: Maximum value for clipping. Defaults to 1.0.
             targeted: Targeted attack if True. Defaults to False.
-            device: Device to use for tensors. Defaults to cuda if available.
         """
 
-        super().__init__(device, normalize)
+        super().__init__(normalize, device)
 
         self.model = model
         self.eps = eps

diff --git a/src/torchattack/geoda.py b/src/torchattack/geoda.py
@@ -96,6 +96,7 @@ def __init__(
         self,
         model: nn.Module,
         normalize: Callable[[torch.Tensor], torch.Tensor] | None,
+        device: torch.device | None = None,
         input_shape: tuple = (3, 224, 224),
         epsilon: int = 5,
         p: str = 'l2',
@@ -107,9 +108,8 @@ def __init__(
         grad_estimator_batch_size: int = 40,
         clip_min: float = 0.0,
         clip_max: float = 1.0,
-        device: torch.device | None = None,
     ):
-        super().__init__(device, normalize)
+        super().__init__(normalize, device)
         self.model = model
 
         self.epsilon = epsilon

diff --git a/src/torchattack/mifgsm.py b/src/torchattack/mifgsm.py
@@ -17,31 +17,31 @@ def __init__(
         self,
         model: nn.Module,
         normalize: Callable[[torch.Tensor], torch.Tensor] | None,
+        device: torch.device | None = None,
         eps: float = 8 / 255,
         steps: int = 10,
         alpha: float | None = None,
         decay: float = 1.0,
         clip_min: float = 0.0,
         clip_max: float = 1.0,
         targeted: bool = False,
-        device: torch.device | None = None,
     ) -> None:
         """Initialize the MI-FGSM attack.
 
         Args:
             model: The model to attack.
             normalize: A transform to normalize images.
+            device: Device to use for tensors. Defaults to cuda if available.
             eps: The maximum perturbation. Defaults to 8/255.
             steps: Number of steps. Defaults to 10.
             alpha: Step size, `eps / steps` if None. Defaults to None.
             decay: Decay factor for the momentum term. Defaults to 1.0.
             clip_min: Minimum value for clipping. Defaults to 0.0.
             clip_max: Maximum value for clipping. Defaults to 1.0.
             targeted: Targeted attack if True. Defaults to False.
-            device: Device to use for tensors. Defaults to cuda if available.
         """
 
-        super().__init__(device, normalize)
+        super().__init__(normalize, device)
 
         self.model = model
         self.eps = eps

diff --git a/src/torchattack/nifgsm.py b/src/torchattack/nifgsm.py
@@ -21,31 +21,31 @@ def __init__(
         self,
         model: nn.Module,
         normalize: Callable[[torch.Tensor], torch.Tensor] | None,
+        device: torch.device | None = None,
         eps: float = 8 / 255,
         steps: int = 10,
         alpha: float | None = None,
         decay: float = 1.0,
         clip_min: float = 0.0,
         clip_max: float = 1.0,
         targeted: bool = False,
-        device: torch.device | None = None,
     ) -> None:
         """Initialize the NI-FGSM attack.
 
         Args:
             model: The model to attack.
             normalize: A transform to normalize images.
+            device: Device to use for tensors. Defaults to cuda if available.
             eps: The maximum perturbation. Defaults to 8/255.
             steps: Number of steps. Defaults to 10.
             alpha: Step size, `eps / steps` if None. Defaults to None.
             decay: Decay factor for the momentum term. Defaults to 1.0.
             clip_min: Minimum value for clipping. Defaults to 0.0.
             clip_max: Maximum value for clipping. Defaults to 1.0.
             targeted: Targeted attack if True. Defaults to False.
-            device: Device to use for tensors. Defaults to cuda if available.
         """
 
-        super().__init__(device, normalize)
+        super().__init__(normalize, device)
 
         self.model = model
         self.eps = eps

diff --git a/src/torchattack/pgd.py b/src/torchattack/pgd.py
@@ -17,31 +17,31 @@ def __init__(
         self,
         model: nn.Module,
         normalize: Callable[[torch.Tensor], torch.Tensor] | None,
+        device: torch.device | None = None,
         eps: float = 8 / 255,
         steps: int = 10,
         alpha: float | None = None,
         random_start: bool = True,
         clip_min: float = 0.0,
         clip_max: float = 1.0,
         targeted: bool = False,
-        device: torch.device | None = None,
     ) -> None:
         """Initialize the PGD attack.
 
         Args:
             model: The model to attack.
             normalize: A transform to normalize images.
+            device: Device to use for tensors. Defaults to cuda if available.
             eps: The maximum perturbation. Defaults to 8/255.
             steps: Number of steps. Defaults to 10.
             alpha: Step size, `eps / steps` if None. Defaults to None.
             random_start: Start from random uniform perturbation. Defaults to True.
             clip_min: Minimum value for clipping. Defaults to 0.0.
             clip_max: Maximum value for clipping. Defaults to 1.0.
             targeted: Targeted attack if True. Defaults to False.
-            device: Device to use for tensors. Defaults to cuda if available.
         """
 
-        super().__init__(device, normalize)
+        super().__init__(normalize, device)
 
         self.model = model
         self.eps = eps

diff --git a/src/torchattack/pgdl2.py b/src/torchattack/pgdl2.py
@@ -19,31 +19,31 @@ def __init__(
         self,
         model: nn.Module,
         normalize: Callable[[torch.Tensor], torch.Tensor] | None,
+        device: torch.device | None = None,
         eps: float = 1.0,
         steps: int = 10,
         alpha: float | None = None,
         random_start: bool = True,
         clip_min: float = 0.0,
         clip_max: float = 1.0,
         targeted: bool = False,
-        device: torch.device | None = None,
     ) -> None:
         """Initialize the PGD attack.
 
         Args:
             model: The model to attack.
             normalize: A transform to normalize images.
+            device: Device to use for tensors. Defaults to cuda if available.
             eps: The maximum perturbation, measured in L2. Defaults to 1.0.
             steps: Number of steps. Defaults to 10.
             alpha: Step size, `eps / steps` if None. Defaults to None.
             random_start: Start from random uniform perturbation. Defaults to True.
             clip_min: Minimum value for clipping. Defaults to 0.0.
             clip_max: Maximum value for clipping. Defaults to 1.0.
             targeted: Targeted attack if True. Defaults to False.
-            device: Device to use for tensors. Defaults to cuda if available.
         """
 
-        super().__init__(device, normalize)
+        super().__init__(normalize, device)
 
         self.model = model
         self.eps = eps

diff --git a/src/torchattack/sinifgsm.py b/src/torchattack/sinifgsm.py
@@ -17,6 +17,7 @@ def __init__(
         self,
         model: nn.Module,
         normalize: Callable[[torch.Tensor], torch.Tensor] | None,
+        device: torch.device | None = None,
         eps: float = 8 / 255,
         steps: int = 10,
         alpha: float | None = None,
@@ -25,13 +26,13 @@ def __init__(
         clip_min: float = 0.0,
         clip_max: float = 1.0,
         targeted: bool = False,
-        device: torch.device | None = None,
     ) -> None:
         """Initialize the SI-NI-FGSM attack.
 
         Args:
             model: The model to attack.
             normalize: A transform to normalize images.
+            device: Device to use for tensors. Defaults to cuda if available.
             eps: The maximum perturbation. Defaults to 8/255.
             steps: Number of steps. Defaults to 10.
             alpha: Step size, `eps / steps` if None. Defaults to None.
@@ -40,10 +41,9 @@ def __init__(
             clip_min: Minimum value for clipping. Defaults to 0.0.
             clip_max: Maximum value for clipping. Defaults to 1.0.
             targeted: Targeted attack if True. Defaults to False.
-            device: Device to use for tensors. Defaults to cuda if available.
         """
 
-        super().__init__(device, normalize)
+        super().__init__(normalize, device)
 
         self.model = model
         self.eps = eps

diff --git a/src/torchattack/tifgsm.py b/src/torchattack/tifgsm.py
@@ -19,6 +19,7 @@ def __init__(
         self,
         model: nn.Module,
         normalize: Callable[[torch.Tensor], torch.Tensor] | None,
+        device: torch.device | None = None,
         eps: float = 8 / 255,
         steps: int = 10,
         alpha: float | None = None,
@@ -28,7 +29,6 @@ def __init__(
         clip_min: float = 0.0,
         clip_max: float = 1.0,
         targeted: bool = False,
-        device: torch.device | None = None,
     ) -> None:
         """Initialize the TI-FGSM attack.
 
@@ -40,6 +40,7 @@ def __init__(
         Args:
             model: The model to attack.
             normalize: A transform to normalize images.
+            device: Device to use for tensors. Defaults to cuda if available.
             eps: The maximum perturbation. Defaults to 8/255.
             steps: Number of steps. Defaults to 10.
             alpha: Step size, `eps / steps` if None. Defaults to None.
@@ -49,10 +50,9 @@ def __init__(
             clip_min: Minimum value for clipping. Defaults to 0.0.
             clip_max: Maximum value for clipping. Defaults to 1.0.
             targeted: Targeted attack if True. Defaults to False.
-            device: Device to use for tensors. Defaults to cuda if available.
         """
 
-        super().__init__(device, normalize)
+        super().__init__(normalize, device)
 
         self.model = model
         self.eps = eps

diff --git a/src/torchattack/vmifgsm.py b/src/torchattack/vmifgsm.py
@@ -18,6 +18,7 @@ def __init__(
         self,
         model: nn.Module,
         normalize: Callable[[torch.Tensor], torch.Tensor] | None,
+        device: torch.device | None = None,
         eps: float = 8 / 255,
         steps: int = 10,
         alpha: float | None = None,
@@ -27,7 +28,6 @@ def __init__(
         clip_min: float = 0.0,
         clip_max: float = 1.0,
         targeted: bool = False,
-        device: torch.device | None = None,
     ) -> None:
         """Initialize the VMI-FGSM attack.
 
@@ -39,6 +39,7 @@ def __init__(
         Args:
             model: The model to attack.
             normalize: A transform to normalize images.
+            device: Device to use for tensors. Defaults to cuda if available.
             eps: The maximum perturbation. Defaults to 8/255.
             steps: Number of steps. Defaults to 10.
             alpha: Step size, `eps / steps` if None. Defaults to None.
@@ -48,10 +49,9 @@ def __init__(
             clip_min: Minimum value for clipping. Defaults to 0.0.
             clip_max: Maximum value for clipping. Defaults to 1.0.
             targeted: Targeted attack if True. Defaults to False.
-            device: Device to use for tensors. Defaults to cuda if available.
         """
 
-        super().__init__(device, normalize)
+        super().__init__(normalize, device)
 
         self.model = model
         self.eps = eps