refactor: rename transform to normalize in attacks, reorder args

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "torchattack"
-version = "0.2.2"
+version = "0.3.0"
 description = "A set of adversarial attacks implemented in PyTorch"
 authors = [{ name = "spencerwooo", email = "[email protected]" }]
 requires-python = ">=3.10,<3.12"

src/torchattack/admix.py

@@ -16,7 +16,7 @@ class Admix(Attack):
     def __init__(
         self,
         model: nn.Module,
-        transform: Callable[[torch.Tensor], torch.Tensor] | None,
+        normalize: Callable[[torch.Tensor], torch.Tensor] | None,
         eps: float = 8 / 255,
         steps: int = 10,
         alpha: float | None = None,
@@ -33,7 +33,7 @@ def __init__(
 
         Args:
             model: The model to attack.
-            transform: A transform to normalize images.
+            normalize: A transform to normalize images.
             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.
@@ -47,7 +47,7 @@ def __init__(
             device: Device to use for tensors. Defaults to cuda if available.
         """
 
-        super().__init__(transform, device)
+        super().__init__(device, normalize)
 
         self.model = model
         self.eps = eps
@@ -91,7 +91,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
             x_admixs = torch.cat([x_admix * scale for scale in scales])
 
             # Compute loss
-            outs = self.model(self.transform(x_admixs))
+            outs = self.model(self.normalize(x_admixs))
 
             # One-hot encode labels for all admixed images
             one_hot = nn.functional.one_hot(y, self.num_classes)

src/torchattack/base.py

@@ -9,17 +9,18 @@ class Attack(ABC):
 
     def __init__(
         self,
-        transform: Callable[[torch.Tensor], torch.Tensor] | None,
         device: torch.device | None,
+        normalize: Callable[[torch.Tensor], torch.Tensor] | None,
     ) -> None:
         super().__init__()
-        # If transform is None, use identity transform.
-        self.transform = transform if transform else lambda x: x
 
-        # Set device to given or defaults to cuda if available.
+        # Set device to given or defaults to cuda if available
         is_cuda = torch.cuda.is_available()
         self.device = device if device else torch.device('cuda' if is_cuda else 'cpu')
 
+        # If normalize is None, use identity function
+        self.normalize = normalize if normalize else lambda x: x
+
     def __call__(self, *args: Any, **kwds: Any) -> Any:
         return self.forward(*args, **kwds)
 
@@ -29,7 +30,7 @@ def __repr__(self) -> str:
         def repr_map(k, v):
             if isinstance(v, float):
                 return f'{k}={v:.3f}'
-            if k in ['model', 'transform']:
+            if k in ['model', 'normalize']:
                 return f'{k}={v.__class__.__name__}'
             if isinstance(v, torch.Tensor):
                 return f'{k}={v.shape}'

src/torchattack/dataset.py

@@ -1,7 +1,6 @@
 import csv
-from typing import Callable
+from typing import Any, Callable
 
-import numpy as np
 import torch
 from PIL import Image
 from torch.utils.data import DataLoader, Dataset
@@ -38,7 +37,7 @@ def __init__(
         Args:
             image_root: Path to the folder containing the images.
             pairs_path: Path to the csv file containing the image names and labels.
-            transform: An optional transform to apply to the images. Defaults to None.
+            normalize: An optional transform to apply to the images. Defaults to None.
             max_samples: Maximum number of samples to load. Defaults to None.
         """
 
@@ -70,13 +69,13 @@ def __init__(
     def __len__(self) -> int:
         return len(self.labels)
 
-    def __getitem__(self, index: int) -> tuple[torch.Tensor, int, str]:
+    def __getitem__(self, index: int) -> tuple[Any, int, str]:
         name = self.names[index]
         label = int(self.labels[index]) - 1
 
         image = Image.open(f'{self.image_root}/{name}.png').convert('RGB')
-        image = np.array(image, dtype=np.uint8)
-        image = torch.from_numpy(image).permute((2, 0, 1)).contiguous().float().div(255)
+        # image = np.array(image, dtype=np.uint8)
+        # image = torch.from_numpy(image).permute((2, 0, 1)).contiguous().float().div(255)
         image = self.transform(image) if self.transform else image
         return image, label, name
 
@@ -108,7 +107,7 @@ def __init__(
         batch_size: int = 1,
         shuffle: bool = False,
         num_workers: int = 4,
-        transform: Callable[[torch.Tensor], torch.Tensor] | None = None,
+        normalize: Callable[[torch.Tensor], torch.Tensor] | None = None,
         max_samples: int | None = None,
     ):
         # Specifing a custom image root directory is useful when evaluating
@@ -118,7 +117,7 @@ def __init__(
             dataset=NIPSDataset(
                 image_root=image_root if image_root else f'{path}/images',
                 pairs_path=pairs_path if pairs_path else f'{path}/images.csv',
-                transform=transform,
+                transform=normalize,
                 max_samples=max_samples,
             ),
             batch_size=batch_size,

src/torchattack/deepfool.py

@@ -16,7 +16,7 @@ class DeepFool(Attack):
     def __init__(
         self,
         model: nn.Module,
-        transform: Callable[[torch.Tensor], torch.Tensor] | None,
+        normalize: Callable[[torch.Tensor], torch.Tensor] | None,
         steps: int = 100,
         overshoot: float = 0.02,
         num_classes: int = 10,
@@ -28,7 +28,7 @@ def __init__(
 
         Args:
             model: The model to attack.
-            transform: A transform to normalize images.
+            normalize: A transform to normalize images.
             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.
@@ -37,7 +37,7 @@ def __init__(
             device: Device to use for tensors. Defaults to cuda if available.
         """
 
-        super().__init__(transform, device)
+        super().__init__(device, normalize)
 
         self.model = model
         self.steps = steps
@@ -58,7 +58,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         """
 
         x.requires_grad_()
-        logits = self.model(self.transform(x))
+        logits = self.model(self.normalize(x))
 
         # Get the classes
         classes = logits.argsort(axis=-1).flip(-1).detach()
@@ -140,7 +140,7 @@ def _get_deltas_logits(
         rows = range(n)
         i0 = classes[:, 0]
 
-        logits = self.model(self.transform(x))
+        logits = self.model(self.normalize(x))
         ik = classes[:, k]
         l0 = logits[rows, i0]
         lk = logits[rows, ik]

src/torchattack/difgsm.py

@@ -17,7 +17,7 @@ class DIFGSM(Attack):
     def __init__(
         self,
         model: nn.Module,
-        transform: Callable[[torch.Tensor], torch.Tensor] | None,
+        normalize: Callable[[torch.Tensor], torch.Tensor] | None,
         eps: float = 8 / 255,
         steps: int = 10,
         alpha: float | None = None,
@@ -39,7 +39,7 @@ def __init__(
 
         Args:
             model: The model to attack.
-            transform: A transform to normalize images.
+            normalize: A transform to normalize images.
             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.
@@ -52,7 +52,7 @@ def __init__(
             device: Device to use for tensors. Defaults to cuda if available.
         """
 
-        super().__init__(transform, device)
+        super().__init__(device, normalize)
 
         self.model = model
         self.eps = eps
@@ -90,7 +90,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
             x_adv = input_diversity(x + delta, self.resize_rate, self.diversity_prob)
 
             # Compute loss
-            outs = self.model(self.transform(x_adv))
+            outs = self.model(self.normalize(x_adv))
             loss = self.lossfn(outs, y)
 
             if self.targeted:

src/torchattack/fgsm.py

@@ -16,7 +16,7 @@ class FGSM(Attack):
     def __init__(
         self,
         model: nn.Module,
-        transform: Callable[[torch.Tensor], torch.Tensor] | None,
+        normalize: Callable[[torch.Tensor], torch.Tensor] | None,
         eps: float = 8 / 255,
         clip_min: float = 0.0,
         clip_max: float = 1.0,
@@ -27,15 +27,15 @@ def __init__(
 
         Args:
             model: A torch.nn.Module network model.
-            transform: A transform to normalize images.
+            normalize: A transform to normalize images.
             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__(transform, device)
+        super().__init__(device, normalize)
 
         self.model = model
         self.eps = eps
@@ -59,7 +59,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         # The original implementation of FGSM is not written in this way.
         delta = torch.zeros_like(x, requires_grad=True)
 
-        outs = self.model(self.transform(x + delta))
+        outs = self.model(self.normalize(x + delta))
         loss = self.lossfn(outs, y)
 
         if self.targeted:

src/torchattack/geoda.py

@@ -95,7 +95,7 @@ class GeoDA(Attack):
     def __init__(
         self,
         model: nn.Module,
-        transform: Callable[[torch.Tensor], torch.Tensor] | None,
+        normalize: Callable[[torch.Tensor], torch.Tensor] | None,
         input_shape: tuple = (3, 224, 224),
         epsilon: int = 5,
         p: str = 'l2',
@@ -109,7 +109,7 @@ def __init__(
         clip_max: float = 1.0,
         device: torch.device | None = None,
     ):
-        super().__init__(transform, device)
+        super().__init__(device, normalize)
         self.model = model
 
         self.epsilon = epsilon
@@ -341,7 +341,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
 
     def predict(self, xs: torch.Tensor) -> torch.Tensor:
         xs = xs.to(self.device)
-        out = self.model(self.transform(xs))
+        out = self.model(self.normalize(xs))
         return out.argmax(dim=1).detach()
 
     def distance(

src/torchattack/mifgsm.py

@@ -16,7 +16,7 @@ class MIFGSM(Attack):
     def __init__(
         self,
         model: nn.Module,
-        transform: Callable[[torch.Tensor], torch.Tensor] | None,
+        normalize: Callable[[torch.Tensor], torch.Tensor] | None,
         eps: float = 8 / 255,
         steps: int = 10,
         alpha: float | None = None,
@@ -30,7 +30,7 @@ def __init__(
 
         Args:
             model: The model to attack.
-            transform: A transform to normalize images.
+            normalize: A transform to normalize images.
             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.
@@ -41,7 +41,7 @@ def __init__(
             device: Device to use for tensors. Defaults to cuda if available.
         """
 
-        super().__init__(transform, device)
+        super().__init__(device, normalize)
 
         self.model = model
         self.eps = eps
@@ -74,7 +74,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         # Perform MI-FGSM
         for _ in range(self.steps):
             # Compute loss
-            outs = self.model(self.transform(x + delta))
+            outs = self.model(self.normalize(x + delta))
             loss = self.lossfn(outs, y)
 
             if self.targeted:

src/torchattack/nifgsm.py

@@ -20,7 +20,7 @@ class NIFGSM(Attack):
     def __init__(
         self,
         model: nn.Module,
-        transform: Callable[[torch.Tensor], torch.Tensor] | None,
+        normalize: Callable[[torch.Tensor], torch.Tensor] | None,
         eps: float = 8 / 255,
         steps: int = 10,
         alpha: float | None = None,
@@ -34,7 +34,7 @@ def __init__(
 
         Args:
             model: The model to attack.
-            transform: A transform to normalize images.
+            normalize: A transform to normalize images.
             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.
@@ -45,7 +45,7 @@ def __init__(
             device: Device to use for tensors. Defaults to cuda if available.
         """
 
-        super().__init__(transform, device)
+        super().__init__(device, normalize)
 
         self.model = model
         self.eps = eps
@@ -82,7 +82,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
             x_nes = x + delta + nes
 
             # Compute loss
-            outs = self.model(self.transform(x_nes))
+            outs = self.model(self.normalize(x_nes))
             loss = self.lossfn(outs, y)
 
             if self.targeted:

src/torchattack/pgd.py

@@ -16,7 +16,7 @@ class PGD(Attack):
     def __init__(
         self,
         model: nn.Module,
-        transform: Callable[[torch.Tensor], torch.Tensor] | None,
+        normalize: Callable[[torch.Tensor], torch.Tensor] | None,
         eps: float = 8 / 255,
         steps: int = 10,
         alpha: float | None = None,
@@ -30,7 +30,7 @@ def __init__(
 
         Args:
             model: The model to attack.
-            transform: A transform to normalize images.
+            normalize: A transform to normalize images.
             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.
@@ -41,7 +41,7 @@ def __init__(
             device: Device to use for tensors. Defaults to cuda if available.
         """
 
-        super().__init__(transform, device)
+        super().__init__(device, normalize)
 
         self.model = model
         self.eps = eps
@@ -80,7 +80,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         # Perform PGD
         for _ in range(self.steps):
             # Compute loss
-            outs = self.model(self.transform(x + delta))
+            outs = self.model(self.normalize(x + delta))
             loss = self.lossfn(outs, y)
 
             if self.targeted: